test_suzieq.py

"""
Tests for SuzieQ Integration - Multi-vendor Drift Detection
Tests network state collection, drift detection, and auto-remediation
"""
import pytest
from datetime import datetime
from agent.suzieq_client import SuzieQClient, DriftDetection
from agent.pipeline_engine import NetworkModel, Device, NetworkIntent, OvergrowthPipeline


@pytest.fixture
def suzieq_client():
    """Create SuzieQ client in mock mode"""
    return SuzieQClient(use_suzieq=True)  # Will use mock mode if suzieq not installed


@pytest.fixture
def sample_network():
    """Create sample network model for drift testing"""
    intent = NetworkIntent(
        description="Test network for SuzieQ drift detection",
        business_requirements=["Multi-vendor support", "Drift detection"],
        constraints=["Budget: $50k", "Timeline: 2 weeks"]
    )
    
    model = NetworkModel(
        name="test_network",
        version="1.0.0",
        intent=intent,
        devices=[
            Device(
                name="spine-01",
                role="spine",
                model="Arista DCS-7280SR-48C6",
                vendor="arista",
                mgmt_ip="10.0.0.1",
                location="DC1",
                interfaces=[
                    {"name": "Ethernet1", "ip": "10.1.1.1/30", "description": "to-leaf-01"}
                ]
            ),
            Device(
                name="leaf-01",
                role="leaf",
                model="Arista DCS-7050SX-64",
                vendor="arista",
                mgmt_ip="10.0.0.11",
                location="DC1",
                interfaces=[
                    {"name": "Ethernet48", "ip": "10.1.1.2/30", "description": "to-spine-01"}
                ]
            )
        ],
        vlans=[
            {"id": 10, "name": "Users"},
            {"id": 20, "name": "Servers"},
            {"id": 99, "name": "Management"}
        ],
        subnets=[
            {"network": "10.10.0.0/16", "gateway": "10.10.0.1", "vlan": 10},
            {"network": "10.20.0.0/16", "gateway": "10.20.0.1", "vlan": 20}
        ],
        routing={"protocol": "bgp", "asn": 65001},
        services=["ntp", "dns", "syslog"]
    )
    
    return model


class TestSuzieQStateCollection:
    """Test network state collection via SuzieQ"""
    
    def test_collect_state_mock_mode(self, suzieq_client):
        """Test state collection in mock mode"""
        devices = [
            {'name': 'spine-01', 'ip': '10.0.0.1', 'username': 'admin', 'password': 'admin'},
            {'name': 'leaf-01', 'ip': '10.0.0.11', 'username': 'admin', 'password': 'admin'}
        ]
        
        result = suzieq_client.collect_network_state(devices)
        
        assert 'devices_polled' in result
        assert result['devices_polled'] == 2
        assert result['mock_mode'] is True
        assert 'collection_time' in result
        
    def test_topology_discovery(self, suzieq_client):
        """Test topology discovery via SuzieQ"""
        topology = suzieq_client.get_topology()
        
        assert 'nodes' in topology
        assert 'edges' in topology
        assert topology['mock_mode'] is True
        
        # Mock topology should have sample data
        assert len(topology['nodes']) > 0
        
    def test_vlan_summary(self, suzieq_client):
        """Test VLAN summary retrieval"""
        vlans = suzieq_client.get_vlan_summary()
        
        assert isinstance(vlans, dict)
        # Mock mode returns device -> vlan list mapping
        assert len(vlans) > 0
        # Check structure - each device should have vlan list
        for device, vlan_list in vlans.items():
            assert isinstance(vlan_list, list)


class TestDriftDetection:
    """Test drift detection and remediation"""
    
    def test_detect_no_drift(self, suzieq_client, sample_network):
        """Test drift detection when network matches SoT"""
        intended = sample_network.to_dict()
        
        drift = suzieq_client.detect_drift(intended)
        
        assert isinstance(drift, DriftDetection)
        assert drift.devices_checked >= 0
        # Mock mode may report some drift for testing
        assert 0.0 <= drift.drift_score <= 1.0
        
    def test_drift_detection_structure(self, suzieq_client, sample_network):
        """Test drift detection returns proper structure"""
        intended = sample_network.to_dict()
        
        drift = suzieq_client.detect_drift(intended)
        
        # Check all drift types are present
        assert hasattr(drift, 'config_mismatches')
        assert hasattr(drift, 'missing_vlans')
        assert hasattr(drift, 'extra_vlans')
        assert hasattr(drift, 'ip_conflicts')
        assert hasattr(drift, 'interface_down')
        assert hasattr(drift, 'routing_issues')
        assert hasattr(drift, 'has_drift')
        assert hasattr(drift, 'drift_score')
        
    def test_drift_to_dict(self, suzieq_client, sample_network):
        """Test drift detection serialization"""
        intended = sample_network.to_dict()
        
        drift = suzieq_client.detect_drift(intended)
        drift_dict = drift.to_dict()
        
        assert 'devices_checked' in drift_dict
        assert 'drift_score' in drift_dict
        assert 'has_drift' in drift_dict
        assert 'config_mismatches' in drift_dict
        assert 'drifts_found' in drift_dict
        
        # Validate data types
        assert isinstance(drift_dict['devices_checked'], int)
        assert isinstance(drift_dict['drift_score'], float)
        assert isinstance(drift_dict['has_drift'], bool)


class TestRemediation:
    """Test remediation plan generation and application"""
    
    def test_generate_remediation_plan(self, suzieq_client, sample_network):
        """Test remediation plan generation from drift"""
        intended = sample_network.to_dict()
        drift = suzieq_client.detect_drift(intended)
        
        remediation = suzieq_client.generate_remediation_plan(drift)
        
        assert isinstance(remediation, list)
        
        # Each remediation item should have proper structure
        for item in remediation:
            assert 'device' in item
            assert 'action' in item  # not issue_type
            assert 'severity' in item
            assert 'auto_fix' in item
            assert item['severity'] in ['critical', 'high', 'medium', 'low', 'ERROR', 'WARNING', 'INFO']
            
    def test_remediation_auto_fix_flags(self, suzieq_client, sample_network):
        """Test that dangerous changes require manual approval"""
        intended = sample_network.to_dict()
        drift = suzieq_client.detect_drift(intended)
        
        remediation = suzieq_client.generate_remediation_plan(drift)
        
        # Check that some items are auto-fix, some require approval
        auto_fix_items = [r for r in remediation if r['auto_fix']]
        manual_items = [r for r in remediation if not r['auto_fix']]
        
        # In mock mode, should have both types
        if len(remediation) > 0:
            # At least validate structure is correct
            for item in remediation:
                if item['severity'] == 'critical':
                    # Critical items might not be auto-fix
                    assert isinstance(item['auto_fix'], bool)
                    
    def test_apply_remediation_auto_only(self, suzieq_client, sample_network):
        """Test applying only auto-approved remediations"""
        intended = sample_network.to_dict()
        drift = suzieq_client.detect_drift(intended)
        remediation = suzieq_client.generate_remediation_plan(drift)
        
        result = suzieq_client.apply_remediation(remediation, auto_approve=True)
        
        assert 'total_actions' in result
        assert 'applied' in result
        assert 'skipped' in result
        assert 'failed' in result
        assert 'actions' in result
        
        # Auto-approve should skip manual items
        assert result['applied'] + result['skipped'] + result['failed'] == result['total_actions']
        
    def test_apply_remediation_manual_approval(self, suzieq_client, sample_network):
        """Test applying remediation with manual approval"""
        intended = sample_network.to_dict()
        drift = suzieq_client.detect_drift(intended)
        remediation = suzieq_client.generate_remediation_plan(drift)
        
        # With auto_approve=False, should skip non-auto-fix items
        result = suzieq_client.apply_remediation(remediation, auto_approve=False)
        
        assert 'total_actions' in result
        assert result['applied'] >= 0
        assert result['skipped'] >= 0
        assert result['total_actions'] == len(remediation)


class TestPipelineIntegration:
    """Test SuzieQ integration with the main pipeline"""
    
    def test_pipeline_has_suzieq(self):
        """Test that pipeline initializes SuzieQ client"""
        pipeline = OvergrowthPipeline()
        
        assert hasattr(pipeline, 'suzieq')
        assert isinstance(pipeline.suzieq, SuzieQClient)
        
    def test_stage7_observability(self, sample_network):
        """Test stage7 observability with SuzieQ"""
        pipeline = OvergrowthPipeline()
        result = pipeline.stage7_observability(sample_network)
        
        assert result['status'] == 'partial'
        assert 'mock_mode' in result
        assert 'collection' in result or 'topology' in result
        
    def test_stage7b_drift_detection(self, sample_network):
        """Test stage7b drift detection"""
        pipeline = OvergrowthPipeline()
        result = pipeline.stage7b_drift_detection(sample_network)
        
        assert 'drift_detected' in result
        assert 'drift_score' in result
        assert 'devices_checked' in result
        assert 'summary' in result
        assert 'mock_mode' in result
        
        # Check summary structure
        summary = result['summary']
        assert 'config_mismatches' in summary
        assert 'missing_vlans' in summary
        assert 'extra_vlans' in summary
        assert 'ip_conflicts' in summary
        assert 'interfaces_down' in summary
        assert 'routing_issues' in summary
        
    def test_stage8_validation_with_drift(self, sample_network):
        """Test stage8 validation with drift detection"""
        pipeline = OvergrowthPipeline()
        result = pipeline.stage8_validation(sample_network)
        
        assert result['status'] == 'completed'
        assert 'validation_passed' in result
        assert 'drift_detection' in result
        assert 'compliance_report' in result
        
        # Check compliance report structure
        compliance = result['compliance_report']
        assert 'network_name' in compliance
        assert 'checked_at' in compliance
        assert 'drift_score' in compliance
        assert 'status' in compliance
        assert compliance['status'] in ['COMPLIANT', 'NON_COMPLIANT']
        
    def test_full_pipeline_with_suzieq(self, sample_network):
        """Test complete pipeline run with SuzieQ integration"""
        pipeline = OvergrowthPipeline()
        
        # Run stages 7, 7b, and 8
        obs_result = pipeline.stage7_observability(sample_network)
        drift_result = pipeline.stage7b_drift_detection(sample_network)
        val_result = pipeline.stage8_validation(sample_network)
        
        # All stages should complete
        assert obs_result['status'] == 'partial'
        assert 'drift_detected' in drift_result
        assert val_result['status'] == 'completed'
        
        # Validation should include drift detection results
        assert val_result['drift_detection'] == drift_result


class TestRealSuzieQConnection:
    """Tests that require actual SuzieQ installation"""
    
    @pytest.mark.skipif(
        SuzieQClient(use_suzieq=True).mock_mode,
        reason="Requires SuzieQ installation"
    )
    def test_real_suzieq_connection(self):
        """Test connection to real SuzieQ instance"""
        client = SuzieQClient(use_suzieq=True)
        
        assert not client.mock_mode
        
        # Try to get topology from real SuzieQ
        topology = client.get_topology()
        
        assert 'nodes' in topology
        assert 'edges' in topology
        assert not topology['mock_mode']


if __name__ == '__main__':
    pytest.main([__file__, '-v'])