llama-stack-mirror/tests/integration/test_xdg_e2e.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

import yaml

from llama_stack.distribution.utils.xdg_utils import (
    get_llama_stack_config_dir,
    get_llama_stack_data_dir,
    get_llama_stack_state_dir,
)


class TestXDGEndToEnd(unittest.TestCase):
    """End-to-end tests for XDG compliance workflows."""

    def setUp(self):
        """Set up test environment."""
        self.original_env = {}
        self.env_vars = [
            "XDG_CONFIG_HOME",
            "XDG_DATA_HOME",
            "XDG_STATE_HOME",
            "XDG_CACHE_HOME",
            "LLAMA_STACK_CONFIG_DIR",
            "SQLITE_STORE_DIR",
            "FILES_STORAGE_DIR",
        ]

        for key in self.env_vars:
            self.original_env[key] = os.environ.get(key)

    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 clear_env_vars(self):
        """Clear all relevant environment variables."""
        for key in self.env_vars:
            os.environ.pop(key, None)

    def create_realistic_legacy_structure(self, base_dir: Path) -> Path:
        """Create a realistic legacy ~/.llama directory structure."""
        legacy_dir = base_dir / ".llama"
        legacy_dir.mkdir()

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

        # Ollama distribution
        ollama_dir = distributions_dir / "ollama"
        ollama_dir.mkdir()

        ollama_run_yaml = ollama_dir / "ollama-run.yaml"
        ollama_run_yaml.write_text("""
version: 2
apis:
  - inference
  - safety
  - memory
  - vector_io
  - agents
  - files
providers:
  inference:
    - provider_type: remote::ollama
      config:
        url: http://localhost:11434
""")

        ollama_build_yaml = ollama_dir / "build.yaml"
        ollama_build_yaml.write_text("""
name: ollama
description: Ollama inference provider
docker_image: ollama:latest
""")

        # Create providers.d structure
        providers_dir = legacy_dir / "providers.d"
        providers_dir.mkdir()

        remote_dir = providers_dir / "remote"
        remote_dir.mkdir()

        inference_dir = remote_dir / "inference"
        inference_dir.mkdir()

        custom_provider = inference_dir / "custom-inference.yaml"
        custom_provider.write_text("""
provider_type: remote::custom
config:
  url: http://localhost:8080
  api_key: test_key
""")

        # Create checkpoints with model files
        checkpoints_dir = legacy_dir / "checkpoints"
        checkpoints_dir.mkdir()

        model_dir = checkpoints_dir / "meta-llama" / "Llama-3.2-1B-Instruct"
        model_dir.mkdir(parents=True)

        # Create fake model files
        (model_dir / "consolidated.00.pth").write_bytes(b"fake model weights" * 1000)
        (model_dir / "params.json").write_text('{"dim": 2048, "n_layers": 22}')
        (model_dir / "tokenizer.model").write_bytes(b"fake tokenizer" * 100)

        # Create runtime with databases
        runtime_dir = legacy_dir / "runtime"
        runtime_dir.mkdir()

        (runtime_dir / "trace_store.db").write_text("SQLite format 3\x00" + "fake database content")
        (runtime_dir / "agent_sessions.db").write_text("SQLite format 3\x00" + "fake agent sessions")

        # Create config files
        (legacy_dir / "config.json").write_text('{"version": "0.2.13", "last_updated": "2024-01-01"}')

        return legacy_dir

    def verify_xdg_migration_complete(self, base_dir: Path, legacy_dir: Path):
        """Verify that migration to XDG structure is complete and correct."""
        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").exists())
        self.assertTrue((config_dir / "distributions" / "ollama" / "ollama-run.yaml").exists())

        # Verify YAML content is preserved
        yaml_content = (config_dir / "distributions" / "ollama" / "ollama-run.yaml").read_text()
        self.assertIn("version: 2", yaml_content)
        self.assertIn("remote::ollama", yaml_content)

        # Verify providers.d moved to config
        self.assertTrue((config_dir / "providers.d").exists())
        self.assertTrue((config_dir / "providers.d" / "remote" / "inference").exists())
        self.assertTrue((config_dir / "providers.d" / "remote" / "inference" / "custom-inference.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").exists())
        self.assertTrue(
            (data_dir / "checkpoints" / "meta-llama" / "Llama-3.2-1B-Instruct" / "consolidated.00.pth").exists()
        )

        # Verify model file content preserved
        model_file = data_dir / "checkpoints" / "meta-llama" / "Llama-3.2-1B-Instruct" / "consolidated.00.pth"
        self.assertGreater(model_file.stat().st_size, 1000)  # Should be substantial size

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

        # Verify database files preserved
        db_file = state_dir / "runtime" / "trace_store.db"
        db_content = db_file.read_text()
        self.assertIn("SQLite format 3", db_content)

    def test_fresh_installation_xdg_compliance(self):
        """Test fresh installation uses XDG-compliant paths."""
        self.clear_env_vars()

        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            # Set custom XDG paths
            os.environ["XDG_CONFIG_HOME"] = str(base_dir / "custom_config")
            os.environ["XDG_DATA_HOME"] = str(base_dir / "custom_data")
            os.environ["XDG_STATE_HOME"] = str(base_dir / "custom_state")

            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
                mock_home.return_value = base_dir
                # Mock that no legacy directory exists
                with patch("llama_stack.distribution.utils.xdg_utils.Path.exists") as mock_exists:
                    mock_exists.return_value = False

                    # Fresh installation should use XDG paths
                    config_dir = get_llama_stack_config_dir()
                    data_dir = get_llama_stack_data_dir()
                    state_dir = get_llama_stack_state_dir()

                    self.assertEqual(config_dir, base_dir / "custom_config" / "llama-stack")
                    self.assertEqual(data_dir, base_dir / "custom_data" / "llama-stack")
                    self.assertEqual(state_dir, base_dir / "custom_state" / "llama-stack")

    def test_complete_migration_workflow(self):
        """Test complete migration workflow from legacy to XDG."""
        self.clear_env_vars()

        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 realistic legacy structure
                legacy_dir = self.create_realistic_legacy_structure(base_dir)

                # Verify legacy structure exists
                self.assertTrue(legacy_dir.exists())
                self.assertTrue((legacy_dir / "distributions" / "ollama" / "ollama-run.yaml").exists())

                # Perform migration
                from llama_stack.cli.migrate_xdg import migrate_to_xdg

                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 migration completed successfully
                self.verify_xdg_migration_complete(base_dir, legacy_dir)

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

    def test_migration_preserves_file_integrity(self):
        """Test that migration preserves file integrity and permissions."""
        self.clear_env_vars()

        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_realistic_legacy_structure(base_dir)

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

                # Perform migration
                from llama_stack.cli.migrate_xdg import migrate_to_xdg

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

                # Verify file integrity
                migrated_config = base_dir / ".config" / "llama-stack" / "distributions" / "ollama" / "ollama-run.yaml"
                self.assertTrue(migrated_config.exists())

                # Verify content is identical
                migrated_content = migrated_config.read_text()
                self.assertIn("version: 2", migrated_content)
                self.assertIn("remote::ollama", migrated_content)

                # Verify permissions preserved
                migrated_stat = migrated_config.stat()
                self.assertEqual(original_stat.st_mode, migrated_stat.st_mode)

    def test_mixed_legacy_and_xdg_environment(self):
        """Test behavior in mixed legacy and XDG environment."""
        self.clear_env_vars()

        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            # Set partial XDG environment
            os.environ["XDG_CONFIG_HOME"] = str(base_dir / "xdg_config")
            # Leave XDG_DATA_HOME and XDG_STATE_HOME unset

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

                # Create legacy directory
                legacy_dir = self.create_realistic_legacy_structure(base_dir)

                # Should use legacy directory since it exists
                config_dir = get_llama_stack_config_dir()
                data_dir = get_llama_stack_data_dir()
                state_dir = get_llama_stack_state_dir()

                self.assertEqual(config_dir, legacy_dir)
                self.assertEqual(data_dir, legacy_dir)
                self.assertEqual(state_dir, legacy_dir)

    def test_template_rendering_with_xdg_paths(self):
        """Test that templates render correctly with XDG paths."""
        self.clear_env_vars()

        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            # Set XDG environment
            os.environ["XDG_STATE_HOME"] = str(base_dir / "state")
            os.environ["XDG_DATA_HOME"] = str(base_dir / "data")

            # Mock shell environment variable expansion
            def mock_env_expand(template_string):
                """Mock shell environment variable expansion."""
                result = template_string
                result = result.replace("${env.XDG_STATE_HOME:-~/.local/state}", str(base_dir / "state"))
                result = result.replace("${env.XDG_DATA_HOME:-~/.local/share}", str(base_dir / "data"))
                return result

            # Test template patterns
            template_patterns = [
                "${env.XDG_STATE_HOME:-~/.local/state}/llama-stack/distributions/ollama",
                "${env.XDG_DATA_HOME:-~/.local/share}/llama-stack/distributions/ollama/files",
            ]

            expected_results = [
                str(base_dir / "state" / "llama-stack" / "distributions" / "ollama"),
                str(base_dir / "data" / "llama-stack" / "distributions" / "ollama" / "files"),
            ]

            for pattern, expected in zip(template_patterns, expected_results, strict=False):
                with self.subTest(pattern=pattern):
                    expanded = mock_env_expand(pattern)
                    self.assertEqual(expanded, expected)

    def test_cli_integration_with_xdg_paths(self):
        """Test CLI integration works correctly with XDG paths."""
        self.clear_env_vars()

        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_realistic_legacy_structure(base_dir)

                # Test CLI migrate command
                import argparse

                from llama_stack.cli.migrate_xdg import MigrateXDG

                parser = argparse.ArgumentParser()
                subparsers = parser.add_subparsers()
                migrate_cmd = MigrateXDG.create(subparsers)

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

                with patch("builtins.print") as mock_print:
                    result = migrate_cmd._run_migrate_xdg_cmd(args)
                    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))

                # Legacy directory should still exist after dry-run
                self.assertTrue(legacy_dir.exists())

    def test_config_dirs_integration_after_migration(self):
        """Test that config_dirs works correctly after migration."""
        self.clear_env_vars()

        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 and migrate legacy structure
                self.create_realistic_legacy_structure(base_dir)

                from llama_stack.cli.migrate_xdg import migrate_to_xdg

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

                # Clear module cache to ensure fresh import
                import sys

                if "llama_stack.distribution.utils.config_dirs" in sys.modules:
                    del sys.modules["llama_stack.distribution.utils.config_dirs"]

                # Import config_dirs after migration
                from llama_stack.distribution.utils.config_dirs import (
                    DEFAULT_CHECKPOINT_DIR,
                    DISTRIBS_BASE_DIR,
                    LLAMA_STACK_CONFIG_DIR,
                    RUNTIME_BASE_DIR,
                )

                # Should use XDG paths
                self.assertEqual(LLAMA_STACK_CONFIG_DIR, base_dir / ".config" / "llama-stack")
                self.assertEqual(DEFAULT_CHECKPOINT_DIR, base_dir / ".local" / "share" / "llama-stack" / "checkpoints")
                self.assertEqual(RUNTIME_BASE_DIR, base_dir / ".local" / "state" / "llama-stack" / "runtime")
                self.assertEqual(DISTRIBS_BASE_DIR, base_dir / ".config" / "llama-stack" / "distributions")

    def test_real_file_operations_with_xdg_paths(self):
        """Test real file operations work correctly with XDG paths."""
        self.clear_env_vars()

        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            # Set XDG environment
            os.environ["XDG_CONFIG_HOME"] = str(base_dir / "config")
            os.environ["XDG_DATA_HOME"] = str(base_dir / "data")
            os.environ["XDG_STATE_HOME"] = str(base_dir / "state")

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

                    # Get XDG paths
                    config_dir = get_llama_stack_config_dir()
                    data_dir = get_llama_stack_data_dir()
                    state_dir = get_llama_stack_state_dir()

                    # Create directories
                    config_dir.mkdir(parents=True)
                    data_dir.mkdir(parents=True)
                    state_dir.mkdir(parents=True)

                    # Test writing configuration files
                    config_file = config_dir / "test_config.yaml"
                    config_data = {"version": "2", "test": True}

                    with open(config_file, "w") as f:
                        yaml.dump(config_data, f)

                    # Test reading configuration files
                    with open(config_file) as f:
                        loaded_config = yaml.safe_load(f)

                    self.assertEqual(loaded_config, config_data)

                    # Test creating nested directory structure
                    model_dir = data_dir / "checkpoints" / "meta-llama" / "test-model"
                    model_dir.mkdir(parents=True)

                    # Test writing large files
                    model_file = model_dir / "model.bin"
                    test_data = b"test model data" * 1000
                    model_file.write_bytes(test_data)

                    # Verify file integrity
                    read_data = model_file.read_bytes()
                    self.assertEqual(read_data, test_data)

                    # Test state files
                    state_file = state_dir / "runtime" / "session.db"
                    state_file.parent.mkdir(parents=True)
                    state_file.write_text("SQLite format 3\x00test database")

                    # Verify state file
                    state_content = state_file.read_text()
                    self.assertIn("SQLite format 3", state_content)

    def test_backwards_compatibility_scenario(self):
        """Test complete backwards compatibility scenario."""
        self.clear_env_vars()

        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            # Scenario: User has existing legacy installation
            legacy_dir = self.create_realistic_legacy_structure(base_dir)

            # User sets legacy environment variable
            os.environ["LLAMA_STACK_CONFIG_DIR"] = str(legacy_dir)

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

                # Should continue using legacy paths
                config_dir = get_llama_stack_config_dir()
                data_dir = get_llama_stack_data_dir()
                state_dir = get_llama_stack_state_dir()

                self.assertEqual(config_dir, legacy_dir)
                self.assertEqual(data_dir, legacy_dir)
                self.assertEqual(state_dir, legacy_dir)

                # Should be able to access existing files
                yaml_file = legacy_dir / "distributions" / "ollama" / "ollama-run.yaml"
                self.assertTrue(yaml_file.exists())

                # Should be able to parse existing configuration
                with open(yaml_file) as f:
                    config = yaml.safe_load(f)

                self.assertEqual(config["version"], 2)

    def test_error_recovery_scenarios(self):
        """Test error recovery in various scenarios."""
        self.clear_env_vars()

        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

                # Scenario 1: Partial migration failure
                self.create_realistic_legacy_structure(base_dir)

                # Create conflicting file in target location
                config_dir = base_dir / ".config" / "llama-stack"
                config_dir.mkdir(parents=True)

                conflicting_file = config_dir / "distributions"
                conflicting_file.touch()  # Create file instead of directory

                from llama_stack.cli.migrate_xdg import migrate_to_xdg

                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)

                        # Should handle conflicts gracefully
                        print_calls = [call[0][0] for call in mock_print.call_args_list]
                        conflict_mentioned = any(
                            "Warning" in call or "conflict" in call.lower() for call in print_calls
                        )

                        # Migration should complete with warnings
                        self.assertTrue(result or conflict_mentioned)

    def test_cross_platform_compatibility(self):
        """Test cross-platform compatibility of XDG implementation."""
        self.clear_env_vars()

        with tempfile.TemporaryDirectory() as temp_dir:
            base_dir = Path(temp_dir)

            # Test with different path separators and formats
            if os.name == "nt":  # Windows
                # Test Windows-style paths
                os.environ["XDG_CONFIG_HOME"] = str(base_dir / "config").replace("/", "\\")
                os.environ["XDG_DATA_HOME"] = str(base_dir / "data").replace("/", "\\")
            else:  # Unix-like
                # Test Unix-style paths
                os.environ["XDG_CONFIG_HOME"] = str(base_dir / "config")
                os.environ["XDG_DATA_HOME"] = str(base_dir / "data")

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

                    # Should work regardless of platform
                    config_dir = get_llama_stack_config_dir()
                    data_dir = get_llama_stack_data_dir()

                    # Paths should be valid for the current platform
                    self.assertTrue(config_dir.is_absolute())
                    self.assertTrue(data_dir.is_absolute())

                    # Should be able to create directories
                    config_dir.mkdir(parents=True, exist_ok=True)
                    data_dir.mkdir(parents=True, exist_ok=True)

                    self.assertTrue(config_dir.exists())
                    self.assertTrue(data_dir.exists())


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