#!/usr/bin/env python3
"""
Unit tests for the Debian APT Stage

This module contains comprehensive tests for the AptStage class to ensure
it correctly handles Debian package management within osbuild.

Author: Debian bootc-image-builder team
License: Same as original bootc-image-builder
"""

import unittest
import tempfile
import os
import json
import shutil
from unittest.mock import Mock, patch, MagicMock
import sys

# Add the osbuild-stages directory to the path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'osbuild-stages', 'apt-stage'))

from apt_stage import AptStage


class TestAptStage(unittest.TestCase):
    """Test cases for the AptStage class."""
    
    def setUp(self):
        """Set up test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        self.test_options = {
            'packages': ['linux-image-amd64', 'systemd', 'initramfs-tools'],
            'release': 'trixie',
            'arch': 'amd64',
            'repos': [
                {
                    'name': 'debian',
                    'url': 'http://deb.debian.org/debian',
                    'suite': 'trixie',
                    'components': ['main', 'contrib']
                }
            ]
        }
        
        # Create a mock context
        self.mock_context = Mock()
        self.mock_context.root = self.temp_dir
        
        # Mock the context.run method
        self.mock_context.run.return_value = Mock(
            returncode=0,
            stdout='',
            stderr=''
        )
    
    def tearDown(self):
        """Clean up test fixtures."""
        shutil.rmtree(self.temp_dir)
    
    def test_initialization(self):
        """Test AptStage initialization with valid options."""
        stage = AptStage(self.test_options)
        
        self.assertEqual(stage.packages, ['linux-image-amd64', 'systemd', 'initramfs-tools'])
        self.assertEqual(stage.release, 'trixie')
        self.assertEqual(stage.arch, 'amd64')
        self.assertEqual(len(stage.repos), 1)
    
    def test_initialization_without_packages(self):
        """Test AptStage initialization fails without packages."""
        options = self.test_options.copy()
        del options['packages']
        
        with self.assertRaises(ValueError) as context:
            AptStage(options)
        
        self.assertIn("No packages specified", str(context.exception))
    
    def test_initialization_with_defaults(self):
        """Test AptStage initialization with minimal options."""
        options = {
            'packages': ['linux-image-amd64']
        }
        
        stage = AptStage(options)
        
        self.assertEqual(stage.release, 'trixie')  # default
        self.assertEqual(stage.arch, 'amd64')      # default
        self.assertTrue(stage.install_weak_deps)   # default
        self.assertEqual(stage.exclude_packages, [])  # default
    
    def test_setup_apt_config(self):
        """Test APT configuration setup."""
        stage = AptStage(self.test_options)
        
        stage._setup_apt_config(self.mock_context)
        
        # Check that the config directory was created
        config_dir = os.path.join(self.temp_dir, 'etc', 'apt', 'apt.conf.d')
        self.assertTrue(os.path.exists(config_dir))
        
        # Check that the config file was created
        config_file = os.path.join(config_dir, '99osbuild')
        self.assertTrue(os.path.exists(config_file))
        
        # Check config file contents
        with open(config_file, 'r') as f:
            config_content = f.read()
        
        # Verify key configuration options are present
        self.assertIn('Acquire::Check-Valid-Until "false"', config_content)
        self.assertIn('Dpkg::Options::="--force-confdef"', config_content)
        self.assertIn('Dpkg::Use-Pty "false"', config_content)
    
    def test_configure_repositories_with_custom_repos(self):
        """Test repository configuration with custom repositories."""
        stage = AptStage(self.test_options)
        
        stage._configure_repositories(self.mock_context)
        
        # Check that the sources directory was created
        sources_dir = os.path.join(self.temp_dir, 'etc', 'apt', 'sources.list.d')
        self.assertTrue(os.path.exists(sources_dir))
        
        # Check that the repository file was created
        repo_file = os.path.join(sources_dir, 'debian.list')
        self.assertTrue(os.path.exists(repo_file))
        
        # Check repository file contents
        with open(repo_file, 'r') as f:
            repo_content = f.read()
        
        expected_content = 'deb http://deb.debian.org/debian trixie main contrib\n'
        self.assertEqual(repo_content, expected_content)
    
    def test_configure_repositories_with_defaults(self):
        """Test repository configuration with default repositories."""
        options = {
            'packages': ['linux-image-amd64'],
            'release': 'bookworm'
        }
        stage = AptStage(options)
        
        stage._configure_repositories(self.mock_context)
        
        # Check that default repositories were created
        sources_dir = os.path.join(self.temp_dir, 'etc', 'apt', 'sources.list.d')
        
        expected_files = ['debian.list', 'debian-security.list', 'debian-updates.list']
        for filename in expected_files:
            filepath = os.path.join(sources_dir, filename)
            self.assertTrue(os.path.exists(filepath))
    
    def test_update_package_lists_success(self):
        """Test successful package list update."""
        stage = AptStage(self.test_options)
        
        stage._update_package_lists(self.mock_context)
        
        # Verify that apt-get update was called
        self.mock_context.run.assert_called_with(['apt-get', 'update'])
    
    def test_update_package_lists_failure(self):
        """Test package list update failure."""
        stage = AptStage(self.test_options)
        
        # Mock a failed command
        self.mock_context.run.return_value = Mock(
            returncode=1,
            stdout='',
            stderr='Failed to update'
        )
        
        with self.assertRaises(RuntimeError) as context:
            stage._update_package_lists(self.mock_context)
        
        self.assertIn("Failed to update package lists", str(context.exception))
    
    def test_install_packages_success(self):
        """Test successful package installation."""
        stage = AptStage(self.test_options)
        
        stage._install_packages(self.mock_context)
        
        # Verify that apt-get install was called with correct arguments
        expected_cmd = [
            'apt-get', 'install', '-y', '--no-install-recommends',
            'linux-image-amd64', 'systemd', 'initramfs-tools'
        ]
        self.mock_context.run.assert_called_with(expected_cmd)
    
    def test_install_packages_with_custom_arch(self):
        """Test package installation with custom architecture."""
        options = self.test_options.copy()
        options['arch'] = 'arm64'
        stage = AptStage(options)
        
        stage._install_packages(self.mock_context)
        
        # Verify that architecture was specified
        expected_cmd = [
            'apt-get', 'install', '-y', '--no-install-recommends',
            '-o', 'APT::Architecture=arm64',
            'linux-image-amd64', 'systemd', 'initramfs-tools'
        ]
        self.mock_context.run.assert_called_with(expected_cmd)
    
    def test_install_packages_failure(self):
        """Test package installation failure."""
        stage = AptStage(self.test_options)
        
        # Mock a failed command
        self.mock_context.run.return_value = Mock(
            returncode=1,
            stdout='',
            stderr='Package installation failed'
        )
        
        with self.assertRaises(RuntimeError) as context:
            stage._install_packages(self.mock_context)
        
        self.assertIn("Package installation failed", str(context.exception))
    
    def test_cleanup_cache_success(self):
        """Test successful cache cleanup."""
        stage = AptStage(self.test_options)
        
        stage._cleanup_cache(self.mock_context)
        
        # Verify that cleanup commands were called
        expected_calls = [
            (['apt-get', 'clean'],),
            (['rm', '-rf', '/var/lib/apt/lists/*'],)
        ]
        
        actual_calls = [call[0] for call in self.mock_context.run.call_args_list]
        self.assertEqual(actual_calls, expected_calls)
    
    def test_cleanup_cache_partial_failure(self):
        """Test cache cleanup with partial failures."""
        stage = AptStage(self.test_options)
        
        # Mock mixed success/failure
        def mock_run(cmd):
            if cmd == ['apt-get', 'clean']:
                return Mock(returncode=1, stderr='Clean failed')
            else:
                return Mock(returncode=0, stderr='')
        
        self.mock_context.run.side_effect = mock_run
        
        # Should not raise an exception, just log warnings
        stage._cleanup_cache(self.mock_context)
    
    def test_log_apt_errors(self):
        """Test APT error logging functionality."""
        stage = AptStage(self.test_options)
        
        # Mock successful commands for error logging
        self.mock_context.run.return_value = Mock(
            returncode=0,
            stdout='No broken packages',
            stderr=''
        )
        
        # Should not raise an exception
        stage._log_apt_errors(self.mock_context)
        
        # Verify that diagnostic commands were called
        expected_calls = [
            (['apt-get', 'check'],),
            (['dpkg', '--audit'],),
            (['dpkg', '-l'],)
        ]
        
        actual_calls = [call[0] for call in self.mock_context.run.call_args_list]
        self.assertEqual(actual_calls, expected_calls)
    
    def test_full_stage_execution(self):
        """Test complete stage execution flow."""
        stage = AptStage(self.test_options)
        
        stage.run(self.mock_context)
        
        # Verify that all major steps were called
        # This is a high-level test to ensure the flow works
        self.assertGreater(self.mock_context.run.call_count, 0)
    
    def test_stage_execution_with_exception(self):
        """Test stage execution handles exceptions properly."""
        stage = AptStage(self.test_options)
        
        # Mock an exception in one of the steps
        self.mock_context.run.side_effect = Exception("Test exception")
        
        with self.assertRaises(Exception) as context:
            stage.run(self.mock_context)
        
        self.assertIn("Test exception", str(context.exception))


class TestAptStageIntegration(unittest.TestCase):
    """Integration tests for AptStage with real filesystem operations."""
    
    def setUp(self):
        """Set up integration test fixtures."""
        self.temp_dir = tempfile.mkdtemp()
        self.test_options = {
            'packages': ['linux-image-amd64'],
            'release': 'trixie',
            'arch': 'amd64'
        }
    
    def tearDown(self):
        """Clean up integration test fixtures."""
        shutil.rmtree(self.temp_dir)
    
    def test_filesystem_operations(self):
        """Test that filesystem operations work correctly."""
        stage = AptStage(self.test_options)
        
        # Create a real context-like object
        class RealContext:
            def __init__(self, root):
                self.root = root
            
            def run(self, cmd):
                # Mock command execution
                return Mock(returncode=0, stdout='', stderr='')
        
        context = RealContext(self.temp_dir)
        
        # Test APT configuration setup
        stage._setup_apt_config(context)
        
        # Verify files were created
        config_file = os.path.join(self.temp_dir, 'etc', 'apt', 'apt.conf.d', '99osbuild')
        self.assertTrue(os.path.exists(config_file))
        
        # Test repository configuration
        stage._configure_repositories(context)
        
        # Verify repository files were created
        sources_dir = os.path.join(self.temp_dir, 'etc', 'apt', 'sources.list.d')
        self.assertTrue(os.path.exists(sources_dir))
        
        # Check that default repositories were created
        expected_files = ['debian.list', 'debian-security.list', 'debian-updates.list']
        for filename in expected_files:
            filepath = os.path.join(sources_dir, filename)
            self.assertTrue(os.path.exists(filepath))


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