mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-07-26 06:07:43 +00:00
489 lines
20 KiB
Python
489 lines
20 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
import argparse
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from io import StringIO
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from llama_stack.cli.migrate_xdg import MigrateXDG, migrate_to_xdg
|
|
|
|
|
|
class TestMigrateXDGCLI(unittest.TestCase):
|
|
"""Test the MigrateXDG CLI command."""
|
|
|
|
def setUp(self):
|
|
"""Set up test environment."""
|
|
self.original_env = {}
|
|
for key in ["XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_STATE_HOME", "LLAMA_STACK_CONFIG_DIR"]:
|
|
self.original_env[key] = os.environ.get(key)
|
|
os.environ.pop(key, None)
|
|
|
|
def tearDown(self):
|
|
"""Clean up test environment."""
|
|
for key, value in self.original_env.items():
|
|
if value is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = value
|
|
|
|
def create_parser_with_migrate_cmd(self):
|
|
"""Create parser with migrate-xdg command."""
|
|
parser = argparse.ArgumentParser()
|
|
subparsers = parser.add_subparsers(dest="command")
|
|
migrate_cmd = MigrateXDG.create(subparsers)
|
|
return parser, migrate_cmd
|
|
|
|
def test_migrate_xdg_command_creation(self):
|
|
"""Test that MigrateXDG command can be created."""
|
|
parser, migrate_cmd = self.create_parser_with_migrate_cmd()
|
|
|
|
self.assertIsInstance(migrate_cmd, MigrateXDG)
|
|
self.assertEqual(migrate_cmd.parser.prog, "llama migrate-xdg")
|
|
self.assertEqual(migrate_cmd.parser.description, "Migrate from legacy ~/.llama to XDG-compliant directories")
|
|
|
|
def test_migrate_xdg_argument_parsing(self):
|
|
"""Test argument parsing for migrate-xdg command."""
|
|
parser, _ = self.create_parser_with_migrate_cmd()
|
|
|
|
# Test with dry-run flag
|
|
args = parser.parse_args(["migrate-xdg", "--dry-run"])
|
|
self.assertEqual(args.command, "migrate-xdg")
|
|
self.assertTrue(args.dry_run)
|
|
|
|
# Test without dry-run flag
|
|
args = parser.parse_args(["migrate-xdg"])
|
|
self.assertEqual(args.command, "migrate-xdg")
|
|
self.assertFalse(args.dry_run)
|
|
|
|
def test_migrate_xdg_help_text(self):
|
|
"""Test help text for migrate-xdg command."""
|
|
parser, _ = self.create_parser_with_migrate_cmd()
|
|
|
|
# Capture help output
|
|
with patch("sys.stdout", new_callable=StringIO) as mock_stdout:
|
|
with patch("sys.exit"):
|
|
try:
|
|
parser.parse_args(["migrate-xdg", "--help"])
|
|
except SystemExit:
|
|
pass
|
|
|
|
help_text = mock_stdout.getvalue()
|
|
self.assertIn("migrate-xdg", help_text)
|
|
self.assertIn("XDG-compliant directories", help_text)
|
|
self.assertIn("--dry-run", help_text)
|
|
|
|
def test_migrate_xdg_command_execution_no_legacy(self):
|
|
"""Test command execution when no legacy directory exists."""
|
|
parser, migrate_cmd = self.create_parser_with_migrate_cmd()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = Path(temp_dir)
|
|
|
|
args = parser.parse_args(["migrate-xdg"])
|
|
|
|
with patch("builtins.print") as mock_print:
|
|
result = migrate_cmd._run_migrate_xdg_cmd(args)
|
|
|
|
# Should succeed when no migration needed
|
|
self.assertEqual(result, 0)
|
|
|
|
# Should print appropriate message
|
|
print_calls = [call[0][0] for call in mock_print.call_args_list]
|
|
self.assertTrue(any("No legacy directory found" in call for call in print_calls))
|
|
|
|
def test_migrate_xdg_command_execution_with_legacy(self):
|
|
"""Test command execution when legacy directory exists."""
|
|
parser, migrate_cmd = self.create_parser_with_migrate_cmd()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = base_dir / ".llama"
|
|
legacy_dir.mkdir()
|
|
(legacy_dir / "test_file").touch()
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
args = parser.parse_args(["migrate-xdg"])
|
|
|
|
with patch("builtins.print") as mock_print:
|
|
result = migrate_cmd._run_migrate_xdg_cmd(args)
|
|
|
|
# Should succeed
|
|
self.assertEqual(result, 0)
|
|
|
|
# Should print migration information
|
|
print_calls = [call[0][0] for call in mock_print.call_args_list]
|
|
self.assertTrue(any("Found legacy directory" in call for call in print_calls))
|
|
|
|
def test_migrate_xdg_command_execution_dry_run(self):
|
|
"""Test command execution with dry-run flag."""
|
|
parser, migrate_cmd = self.create_parser_with_migrate_cmd()
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = base_dir / ".llama"
|
|
legacy_dir.mkdir()
|
|
(legacy_dir / "test_file").touch()
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
args = parser.parse_args(["migrate-xdg", "--dry-run"])
|
|
|
|
with patch("builtins.print") as mock_print:
|
|
result = migrate_cmd._run_migrate_xdg_cmd(args)
|
|
|
|
# Should succeed
|
|
self.assertEqual(result, 0)
|
|
|
|
# Should print dry-run information
|
|
print_calls = [call[0][0] for call in mock_print.call_args_list]
|
|
self.assertTrue(any("Dry run mode" in call for call in print_calls))
|
|
|
|
def test_migrate_xdg_command_execution_error_handling(self):
|
|
"""Test command execution with error handling."""
|
|
parser, migrate_cmd = self.create_parser_with_migrate_cmd()
|
|
|
|
args = parser.parse_args(["migrate-xdg"])
|
|
|
|
# Mock migrate_to_xdg to raise an exception
|
|
with patch("llama_stack.cli.migrate_xdg.migrate_to_xdg") as mock_migrate:
|
|
mock_migrate.side_effect = Exception("Test error")
|
|
|
|
with patch("builtins.print") as mock_print:
|
|
result = migrate_cmd._run_migrate_xdg_cmd(args)
|
|
|
|
# Should return error code
|
|
self.assertEqual(result, 1)
|
|
|
|
# Should print error message
|
|
print_calls = [call[0][0] for call in mock_print.call_args_list]
|
|
self.assertTrue(any("Error during migration" in call for call in print_calls))
|
|
|
|
def test_migrate_xdg_command_integration(self):
|
|
"""Test full integration of migrate-xdg command."""
|
|
from llama_stack.cli.llama import LlamaCLIParser
|
|
|
|
# Create main parser
|
|
main_parser = LlamaCLIParser()
|
|
|
|
# Test that migrate-xdg is in the subcommands
|
|
with patch("sys.argv", ["llama", "migrate-xdg", "--help"]):
|
|
with patch("sys.stdout", new_callable=StringIO) as mock_stdout:
|
|
with patch("sys.exit"):
|
|
try:
|
|
main_parser.parse_args()
|
|
except SystemExit:
|
|
pass
|
|
|
|
help_text = mock_stdout.getvalue()
|
|
self.assertIn("migrate-xdg", help_text)
|
|
|
|
|
|
class TestMigrateXDGFunction(unittest.TestCase):
|
|
"""Test the migrate_to_xdg function directly."""
|
|
|
|
def setUp(self):
|
|
"""Set up test environment."""
|
|
self.original_env = {}
|
|
for key in ["XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_STATE_HOME", "LLAMA_STACK_CONFIG_DIR"]:
|
|
self.original_env[key] = os.environ.get(key)
|
|
os.environ.pop(key, None)
|
|
|
|
def tearDown(self):
|
|
"""Clean up test environment."""
|
|
for key, value in self.original_env.items():
|
|
if value is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = value
|
|
|
|
def create_legacy_structure(self, base_dir: Path) -> Path:
|
|
"""Create a test legacy directory structure."""
|
|
legacy_dir = base_dir / ".llama"
|
|
legacy_dir.mkdir()
|
|
|
|
# Create distributions
|
|
(legacy_dir / "distributions").mkdir()
|
|
(legacy_dir / "distributions" / "ollama").mkdir()
|
|
(legacy_dir / "distributions" / "ollama" / "run.yaml").write_text("version: 2\n")
|
|
|
|
# Create checkpoints
|
|
(legacy_dir / "checkpoints").mkdir()
|
|
(legacy_dir / "checkpoints" / "model.bin").write_text("fake model")
|
|
|
|
# Create providers.d
|
|
(legacy_dir / "providers.d").mkdir()
|
|
(legacy_dir / "providers.d" / "provider.yaml").write_text("provider: test\n")
|
|
|
|
# Create runtime
|
|
(legacy_dir / "runtime").mkdir()
|
|
(legacy_dir / "runtime" / "trace.db").write_text("fake database")
|
|
|
|
return legacy_dir
|
|
|
|
def test_migrate_to_xdg_no_legacy_directory(self):
|
|
"""Test migrate_to_xdg when no legacy directory exists."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = Path(temp_dir)
|
|
|
|
result = migrate_to_xdg(dry_run=False)
|
|
self.assertTrue(result)
|
|
|
|
def test_migrate_to_xdg_dry_run(self):
|
|
"""Test migrate_to_xdg with dry_run=True."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = self.create_legacy_structure(base_dir)
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
with patch("builtins.print") as mock_print:
|
|
result = migrate_to_xdg(dry_run=True)
|
|
self.assertTrue(result)
|
|
|
|
# Should print dry run information
|
|
print_calls = [call[0][0] for call in mock_print.call_args_list]
|
|
self.assertTrue(any("Dry run mode" in call for call in print_calls))
|
|
|
|
# Legacy directory should still exist
|
|
self.assertTrue(legacy_dir.exists())
|
|
|
|
def test_migrate_to_xdg_user_confirms(self):
|
|
"""Test migrate_to_xdg when user confirms migration."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = self.create_legacy_structure(base_dir)
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
with patch("builtins.input") as mock_input:
|
|
mock_input.side_effect = ["y", "y"] # Confirm migration and cleanup
|
|
|
|
result = migrate_to_xdg(dry_run=False)
|
|
self.assertTrue(result)
|
|
|
|
# Legacy directory should be removed
|
|
self.assertFalse(legacy_dir.exists())
|
|
|
|
# XDG directories should be created
|
|
self.assertTrue((base_dir / ".config" / "llama-stack").exists())
|
|
self.assertTrue((base_dir / ".local" / "share" / "llama-stack").exists())
|
|
|
|
def test_migrate_to_xdg_user_cancels(self):
|
|
"""Test migrate_to_xdg when user cancels migration."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = self.create_legacy_structure(base_dir)
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
with patch("builtins.input") as mock_input:
|
|
mock_input.return_value = "n" # Cancel migration
|
|
|
|
result = migrate_to_xdg(dry_run=False)
|
|
self.assertFalse(result)
|
|
|
|
# Legacy directory should still exist
|
|
self.assertTrue(legacy_dir.exists())
|
|
|
|
def test_migrate_to_xdg_partial_migration(self):
|
|
"""Test migrate_to_xdg with partial migration (some files fail)."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = self.create_legacy_structure(base_dir)
|
|
self.assertFalse(legacy_dir.exists())
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
# Create target directory with conflicting file
|
|
config_dir = base_dir / ".config" / "llama-stack"
|
|
config_dir.mkdir(parents=True)
|
|
(config_dir / "distributions").mkdir()
|
|
(config_dir / "distributions" / "existing.yaml").write_text("existing")
|
|
|
|
with patch("builtins.input") as mock_input:
|
|
mock_input.side_effect = ["y", "n"] # Confirm migration, don't cleanup
|
|
|
|
with patch("builtins.print") as mock_print:
|
|
result = migrate_to_xdg(dry_run=False)
|
|
self.assertTrue(result)
|
|
|
|
# Should print warning about conflicts
|
|
print_calls = [call[0][0] for call in mock_print.call_args_list]
|
|
self.assertTrue(any("Warning: Target already exists" in call for call in print_calls))
|
|
|
|
def test_migrate_to_xdg_permission_error(self):
|
|
"""Test migrate_to_xdg with permission errors."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = self.create_legacy_structure(base_dir)
|
|
self.assertFalse(legacy_dir.exists())
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
# Create readonly target directory
|
|
config_dir = base_dir / ".config" / "llama-stack"
|
|
config_dir.mkdir(parents=True)
|
|
config_dir.chmod(0o444) # Read-only
|
|
|
|
try:
|
|
with patch("builtins.input") as mock_input:
|
|
mock_input.side_effect = ["y", "n"] # Confirm migration, don't cleanup
|
|
|
|
with patch("builtins.print") as mock_print:
|
|
result = migrate_to_xdg(dry_run=False)
|
|
self.assertFalse(result)
|
|
|
|
# Should handle permission errors gracefully
|
|
print_calls = [call[0][0] for call in mock_print.call_args_list]
|
|
# Should contain some error or warning message
|
|
self.assertTrue(len(print_calls) > 0)
|
|
|
|
finally:
|
|
# Restore permissions for cleanup
|
|
config_dir.chmod(0o755)
|
|
|
|
def test_migrate_to_xdg_empty_legacy_directory(self):
|
|
"""Test migrate_to_xdg with empty legacy directory."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = base_dir / ".llama"
|
|
legacy_dir.mkdir() # Empty directory
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
result = migrate_to_xdg(dry_run=False)
|
|
self.assertTrue(result)
|
|
|
|
def test_migrate_to_xdg_preserves_file_content(self):
|
|
"""Test that migrate_to_xdg preserves file content correctly."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = self.create_legacy_structure(base_dir)
|
|
|
|
# Add specific content to test
|
|
test_content = "test configuration content"
|
|
(legacy_dir / "distributions" / "ollama" / "run.yaml").write_text(test_content)
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
with patch("builtins.input") as mock_input:
|
|
mock_input.side_effect = ["y", "y"] # Confirm migration and cleanup
|
|
|
|
result = migrate_to_xdg(dry_run=False)
|
|
self.assertTrue(result)
|
|
|
|
# Check content was preserved
|
|
migrated_file = base_dir / ".config" / "llama-stack" / "distributions" / "ollama" / "run.yaml"
|
|
self.assertTrue(migrated_file.exists())
|
|
self.assertEqual(migrated_file.read_text(), test_content)
|
|
|
|
def test_migrate_to_xdg_with_symlinks(self):
|
|
"""Test migrate_to_xdg with symbolic links."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = self.create_legacy_structure(base_dir)
|
|
|
|
# Create symlink
|
|
actual_file = legacy_dir / "actual_config.yaml"
|
|
actual_file.write_text("actual config")
|
|
|
|
symlink_file = legacy_dir / "distributions" / "symlinked.yaml"
|
|
symlink_file.symlink_to(actual_file)
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
with patch("builtins.input") as mock_input:
|
|
mock_input.side_effect = ["y", "y"] # Confirm migration and cleanup
|
|
|
|
result = migrate_to_xdg(dry_run=False)
|
|
self.assertTrue(result)
|
|
|
|
# Check symlink was preserved
|
|
migrated_symlink = base_dir / ".config" / "llama-stack" / "distributions" / "symlinked.yaml"
|
|
self.assertTrue(migrated_symlink.exists())
|
|
self.assertTrue(migrated_symlink.is_symlink())
|
|
|
|
def test_migrate_to_xdg_nested_directory_structure(self):
|
|
"""Test migrate_to_xdg with nested directory structures."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = self.create_legacy_structure(base_dir)
|
|
|
|
# Create nested structure
|
|
nested_dir = legacy_dir / "checkpoints" / "org" / "model" / "variant"
|
|
nested_dir.mkdir(parents=True)
|
|
(nested_dir / "model.bin").write_text("nested model")
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
with patch("builtins.input") as mock_input:
|
|
mock_input.side_effect = ["y", "y"] # Confirm migration and cleanup
|
|
|
|
result = migrate_to_xdg(dry_run=False)
|
|
self.assertTrue(result)
|
|
|
|
# Check nested structure was preserved
|
|
migrated_nested = (
|
|
base_dir / ".local" / "share" / "llama-stack" / "checkpoints" / "org" / "model" / "variant"
|
|
)
|
|
self.assertTrue(migrated_nested.exists())
|
|
self.assertTrue((migrated_nested / "model.bin").exists())
|
|
|
|
def test_migrate_to_xdg_user_input_variations(self):
|
|
"""Test migrate_to_xdg with various user input variations."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
base_dir = Path(temp_dir)
|
|
legacy_dir = self.create_legacy_structure(base_dir)
|
|
|
|
with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
|
|
mock_home.return_value = base_dir
|
|
|
|
# Test various forms of "yes"
|
|
for yes_input in ["y", "Y", "yes", "Yes", "YES"]:
|
|
# Recreate legacy directory for each test
|
|
if legacy_dir.exists():
|
|
import shutil
|
|
|
|
shutil.rmtree(legacy_dir)
|
|
self.create_legacy_structure(base_dir)
|
|
|
|
with patch("builtins.input") as mock_input:
|
|
mock_input.side_effect = [yes_input, "n"] # Confirm migration, don't cleanup
|
|
|
|
result = migrate_to_xdg(dry_run=False)
|
|
self.assertTrue(result, f"Failed with input: {yes_input}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|