llama-stack-mirror/tests/integration/test_xdg_migration.py

# 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 os
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from llama_stack.cli.migrate_xdg import migrate_to_xdg


class TestXDGMigrationIntegration(unittest.TestCase):
    """Integration tests for XDG migration functionality."""

    def setUp(self):
        """Set up test environment."""
        # Store original environment variables
        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)

        # Clear environment variables
        for key in self.original_env:
            os.environ.pop(key, None)

    def tearDown(self):
        """Clean up test environment."""
        # Restore original environment variables
        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 realistic legacy ~/.llama directory structure."""
        legacy_dir = base_dir / ".llama"
        legacy_dir.mkdir()

        # Create distributions
        distributions_dir = legacy_dir / "distributions"
        distributions_dir.mkdir()

        # Create sample distribution
        ollama_dir = distributions_dir / "ollama"
        ollama_dir.mkdir()
        (ollama_dir / "ollama-run.yaml").write_text("version: 2\napis: []\n")
        (ollama_dir / "build.yaml").write_text("name: ollama\n")

        # Create providers.d
        providers_dir = legacy_dir / "providers.d"
        providers_dir.mkdir()
        (providers_dir / "remote").mkdir()
        (providers_dir / "remote" / "inference").mkdir()
        (providers_dir / "remote" / "inference" / "custom.yaml").write_text("provider_type: remote::custom\n")

        # Create checkpoints
        checkpoints_dir = legacy_dir / "checkpoints"
        checkpoints_dir.mkdir()
        model_dir = checkpoints_dir / "meta-llama" / "Llama-3.2-1B-Instruct"
        model_dir.mkdir(parents=True)
        (model_dir / "consolidated.00.pth").write_text("fake model weights")
        (model_dir / "params.json").write_text('{"dim": 2048}')

        # Create runtime
        runtime_dir = legacy_dir / "runtime"
        runtime_dir.mkdir()
        (runtime_dir / "trace_store.db").write_text("fake sqlite database")

        # Create some fake files in various subdirectories
        (legacy_dir / "config.json").write_text('{"version": "0.2.13"}')

        return legacy_dir

    def verify_xdg_structure(self, base_dir: Path, legacy_dir: Path):
        """Verify that the XDG structure was created correctly."""
        config_dir = base_dir / ".config" / "llama-stack"
        data_dir = base_dir / ".local" / "share" / "llama-stack"
        state_dir = base_dir / ".local" / "state" / "llama-stack"

        # Verify distributions moved to config
        self.assertTrue((config_dir / "distributions").exists())
        self.assertTrue((config_dir / "distributions" / "ollama" / "ollama-run.yaml").exists())

        # Verify providers.d moved to config
        self.assertTrue((config_dir / "providers.d").exists())
        self.assertTrue((config_dir / "providers.d" / "remote" / "inference" / "custom.yaml").exists())

        # Verify checkpoints moved to data
        self.assertTrue((data_dir / "checkpoints").exists())
        self.assertTrue(
            (data_dir / "checkpoints" / "meta-llama" / "Llama-3.2-1B-Instruct" / "consolidated.00.pth").exists()
        )

        # Verify runtime moved to state
        self.assertTrue((state_dir / "runtime").exists())
        self.assertTrue((state_dir / "runtime" / "trace_store.db").exists())

    def test_full_migration_workflow(self):
        """Test complete migration workflow from legacy to XDG."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            # Set up fake home directory
            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create legacy structure
                legacy_dir = self.create_legacy_structure(base_dir)

                # Verify legacy structure exists
                self.assertTrue(legacy_dir.exists())
                self.assertTrue((legacy_dir / "distributions").exists())
                self.assertTrue((legacy_dir / "checkpoints").exists())

                # Perform migration with user confirmation
                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)

                # Verify XDG structure was created
                self.verify_xdg_structure(base_dir, legacy_dir)

                # Verify legacy directory was removed
                self.assertFalse(legacy_dir.exists())

    def test_migration_dry_run(self):
        """Test dry run migration (no actual file movement)."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create legacy structure
                legacy_dir = self.create_legacy_structure(base_dir)

                # Perform dry run migration
                with patch("builtins.print") as mock_print:
                    result = migrate_to_xdg(dry_run=True)
                    self.assertTrue(result)

                    # Check that dry run message was printed
                    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))

                # Verify nothing was actually moved
                self.assertTrue(legacy_dir.exists())
                self.assertTrue((legacy_dir / "distributions").exists())
                self.assertFalse((base_dir / ".config" / "llama-stack").exists())

    def test_migration_user_cancellation(self):
        """Test migration when user cancels."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create legacy structure
                legacy_dir = self.create_legacy_structure(base_dir)

                # User cancels migration
                with patch("builtins.input") as mock_input:
                    mock_input.return_value = "n"

                    result = migrate_to_xdg(dry_run=False)
                    self.assertFalse(result)

                # Verify nothing was moved
                self.assertTrue(legacy_dir.exists())
                self.assertTrue((legacy_dir / "distributions").exists())
                self.assertFalse((base_dir / ".config" / "llama-stack").exists())

    def test_migration_with_existing_xdg_directories(self):
        """Test migration when XDG directories already exist."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create legacy structure
                legacy_dir = self.create_legacy_structure(base_dir)

                # Create existing XDG structure with conflicting files
                config_dir = base_dir / ".config" / "llama-stack"
                config_dir.mkdir(parents=True)
                (config_dir / "distributions").mkdir()
                (config_dir / "distributions" / "existing.yaml").write_text("existing config")

                # Perform migration
                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)

                        # Check that warning was printed
                        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))

                # Verify existing file wasn't overwritten
                self.assertTrue((config_dir / "distributions" / "existing.yaml").exists())

                # Legacy distributions should still exist due to conflict
                self.assertTrue((legacy_dir / "distributions").exists())

    def test_migration_partial_success(self):
        """Test migration when some items succeed and others fail."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create legacy structure
                self.create_legacy_structure(base_dir)

                # Create readonly target directory to simulate permission error
                config_dir = base_dir / ".config" / "llama-stack"
                config_dir.mkdir(parents=True)
                distributions_target = config_dir / "distributions"
                distributions_target.mkdir()
                distributions_target.chmod(0o444)  # Read-only

                try:
                    # Perform migration
                    with patch("builtins.input") as mock_input:
                        mock_input.side_effect = ["y", "n"]  # Confirm migration, don't cleanup

                        migrate_to_xdg(dry_run=False)
                        # Should return True even with partial success

                        # Some items should have been migrated successfully
                        self.assertTrue((base_dir / ".local" / "share" / "llama-stack" / "checkpoints").exists())

                finally:
                    # Restore permissions for cleanup
                    distributions_target.chmod(0o755)

    def test_migration_empty_legacy_directory(self):
        """Test migration when legacy directory exists but is empty."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create empty legacy directory
                legacy_dir = base_dir / ".llama"
                legacy_dir.mkdir()

                result = migrate_to_xdg(dry_run=False)
                self.assertTrue(result)

    def test_migration_preserves_file_permissions(self):
        """Test that migration preserves file permissions."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create legacy structure with specific permissions
                legacy_dir = self.create_legacy_structure(base_dir)

                # Set specific permissions on a file
                config_file = legacy_dir / "distributions" / "ollama" / "ollama-run.yaml"
                config_file.chmod(0o600)
                original_stat = config_file.stat()

                # Perform migration
                with patch("builtins.input") as mock_input:
                    mock_input.side_effect = ["y", "y"]
                    result = migrate_to_xdg(dry_run=False)
                    self.assertTrue(result)

                # Verify permissions were preserved
                migrated_file = base_dir / ".config" / "llama-stack" / "distributions" / "ollama" / "ollama-run.yaml"
                self.assertTrue(migrated_file.exists())
                migrated_stat = migrated_file.stat()
                self.assertEqual(original_stat.st_mode, migrated_stat.st_mode)

    def test_migration_preserves_directory_structure(self):
        """Test that migration preserves complex directory structures."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create complex legacy structure
                legacy_dir = base_dir / ".llama"
                legacy_dir.mkdir()

                # Create nested structure
                complex_path = legacy_dir / "checkpoints" / "org" / "model" / "variant" / "files"
                complex_path.mkdir(parents=True)
                (complex_path / "model.bin").write_text("model data")
                (complex_path / "config.json").write_text("config data")

                # Perform migration
                with patch("builtins.input") as mock_input:
                    mock_input.side_effect = ["y", "y"]
                    result = migrate_to_xdg(dry_run=False)
                    self.assertTrue(result)

                # Verify structure was preserved
                migrated_path = (
                    base_dir
                    / ".local"
                    / "share"
                    / "llama-stack"
                    / "checkpoints"
                    / "org"
                    / "model"
                    / "variant"
                    / "files"
                )
                self.assertTrue(migrated_path.exists())
                self.assertTrue((migrated_path / "model.bin").exists())
                self.assertTrue((migrated_path / "config.json").exists())

    def test_migration_with_symlinks(self):
        """Test migration with symbolic links in legacy directory."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create legacy structure
                legacy_dir = self.create_legacy_structure(base_dir)

                # Create a symlink
                actual_file = legacy_dir / "actual_config.yaml"
                actual_file.write_text("actual config content")

                symlink_file = legacy_dir / "distributions" / "symlinked_config.yaml"
                symlink_file.symlink_to(actual_file)

                # Perform migration
                with patch("builtins.input") as mock_input:
                    mock_input.side_effect = ["y", "y"]
                    result = migrate_to_xdg(dry_run=False)
                    self.assertTrue(result)

                # Verify symlink was preserved
                migrated_symlink = base_dir / ".config" / "llama-stack" / "distributions" / "symlinked_config.yaml"
                self.assertTrue(migrated_symlink.exists())
                self.assertTrue(migrated_symlink.is_symlink())

    def test_migration_large_files(self):
        """Test migration with large files (simulated)."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create legacy structure
                legacy_dir = self.create_legacy_structure(base_dir)

                # Create a larger file (1MB)
                large_file = legacy_dir / "checkpoints" / "large_model.bin"
                large_file.write_bytes(b"0" * (1024 * 1024))

                # Perform migration
                with patch("builtins.input") as mock_input:
                    mock_input.side_effect = ["y", "y"]
                    result = migrate_to_xdg(dry_run=False)
                    self.assertTrue(result)

                # Verify large file was moved correctly
                migrated_file = base_dir / ".local" / "share" / "llama-stack" / "checkpoints" / "large_model.bin"
                self.assertTrue(migrated_file.exists())
                self.assertEqual(migrated_file.stat().st_size, 1024 * 1024)

    def test_migration_with_unicode_filenames(self):
        """Test migration with unicode filenames."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create legacy structure with unicode filenames
                legacy_dir = base_dir / ".llama"
                legacy_dir.mkdir()

                unicode_dir = legacy_dir / "distributions" / "配置"
                unicode_dir.mkdir(parents=True)
                unicode_file = unicode_dir / "模型.yaml"
                unicode_file.write_text("unicode content")

                # Perform migration
                with patch("builtins.input") as mock_input:
                    mock_input.side_effect = ["y", "y"]
                    result = migrate_to_xdg(dry_run=False)
                    self.assertTrue(result)

                # Verify unicode files were migrated
                migrated_file = base_dir / ".config" / "llama-stack" / "distributions" / "配置" / "模型.yaml"
                self.assertTrue(migrated_file.exists())
                self.assertEqual(migrated_file.read_text(), "unicode content")


class TestXDGMigrationCLI(unittest.TestCase):
    """Test the CLI interface for XDG migration."""

    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 test_cli_migrate_command_exists(self):
        """Test that the migrate-xdg CLI command is properly registered."""
        from llama_stack.cli.llama import LlamaCLIParser

        parser = LlamaCLIParser()

        # Parse help to see if migrate-xdg is listed
        with patch("sys.argv", ["llama", "--help"]):
            with patch("sys.exit"):
                with patch("builtins.print") as mock_print:
                    try:
                        parser.parse_args()
                    except SystemExit:
                        pass

                    # Check if migrate-xdg appears in help output
                    help_output = "\n".join([call[0][0] for call in mock_print.call_args_list])
                    self.assertIn("migrate-xdg", help_output)

    def test_cli_migrate_dry_run(self):
        """Test CLI migrate command with dry-run flag."""
        import argparse

        from llama_stack.cli.migrate_xdg import MigrateXDG

        # Create parser and add migrate command
        parser = argparse.ArgumentParser()
        subparsers = parser.add_subparsers()
        MigrateXDG.create(subparsers)

        # Test dry-run flag
        args = parser.parse_args(["migrate-xdg", "--dry-run"])
        self.assertTrue(args.dry_run)

        # Test without dry-run flag
        args = parser.parse_args(["migrate-xdg"])
        self.assertFalse(args.dry_run)

    def test_cli_migrate_execution(self):
        """Test CLI migrate command execution."""
        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir

                # Create legacy directory
                legacy_dir = base_dir / ".llama"
                legacy_dir.mkdir()
                (legacy_dir / "test_file").touch()

                import argparse

                from llama_stack.cli.migrate_xdg import MigrateXDG

                # Create parser and command
                parser = argparse.ArgumentParser()
                subparsers = parser.add_subparsers()
                migrate_cmd = MigrateXDG.create(subparsers)

                # Parse arguments
                args = parser.parse_args(["migrate-xdg", "--dry-run"])

                # Execute command
                with patch("builtins.print") as mock_print:
                    migrate_cmd._run_migrate_xdg_cmd(args)

                    # Verify output was printed
                    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))


if __name__ == "__main__":
    unittest.main()