from unittest.mock import patch, MagicMock
import json
import pytest
from armorer.scanner import VulnerabilityScanner

def test_scanner_generic_high_severity():
    scanner = VulnerabilityScanner()
    
    with patch("subprocess.run") as mock_run:
        # 1. Safety Output (Success)
        # 2. Bandit Output (Generic High Severity)
        bandit_output = json.dumps({
            "metrics": {
                "_totals": {
                    "SEVERITY.HIGH": 1,
                    "SEVERITY.MEDIUM": 0
                }
            }
        })
        
        def side_effect(cmd, **kwargs):
            mock = MagicMock()
            if "safety" in cmd:
                mock.returncode = 0
            elif "bandit" in cmd:
                mock.stdout = bandit_output
                mock.returncode = 0
            return mock
            
        mock_run.side_effect = side_effect
        
        result = scanner.scan()
        
        assert result.agent_risks_found is False
        assert result.critical_found is True
        assert len(result.issues) > 0
        assert "Found 1 HIGH severity" in result.issues[0]

def test_scanner_agent_specific_risk():
    scanner = VulnerabilityScanner()
    
    with patch("subprocess.run") as mock_run:
        # Bandit Output (Agent Risk B602)
        bandit_agent_risk = json.dumps({
            "results": [
                {
                    "test_id": "B602", 
                    "issue_text": "subprocess call with shell=True identified", 
                    "filename": "agent.py", 
                    "line_number": 42
                }
            ],
            "metrics": {
                "_totals": { "SEVERITY.HIGH": 1 }
            }
        })

        def side_effect(cmd, **kwargs):
            mock = MagicMock()
            if "safety" in cmd: mock.returncode = 0
            elif "bandit" in cmd:
                mock.stdout = bandit_agent_risk
                mock.returncode = 0
            return mock
            
        mock_run.side_effect = side_effect
        
        result = scanner.scan()
        
        assert result.agent_risks_found is True
        assert result.critical_found is True
        assert "AGENT RISK" in result.issues[0]
