a2a-source-edit: write tests/test_agent.py

tests/test_agent.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import pytest
 
-from a2a_pack import LocalRunContext, LocalWorkspaceClient, NoAuth, WorkspaceMode
+from a2a_pack import LocalRunContext, LocalWorkspaceClient, NoAuth
 
 from agent import (
     AgentComplianceAuditor,
@@ -40,16 +40,16 @@ async def test_clean_synthetic_audit_outputs_expected_files() -> None:
     options = AuditOptions(audit_id="clean", probe_mode="none", max_files=10)
 
     inv = await agent.inventory_agent(ctx, InventoryRequest(target=target, options=options))
-    cfg = await agent.audit_configuration(ctx, AuditConfigurationRequest(target=target, options=options, inventory=inv.inventory))
-    rt = await agent.audit_runtime(ctx, AuditRuntimeRequest(target=target, options=options, inventory=inv.inventory))
-    risk = await agent.assess_risk(ctx, AssessRiskRequest(target=target, options=options, configuration_findings=cfg.findings, runtime_findings=rt.findings))
-    rem = await agent.generate_remediation(ctx, GenerateRemediationRequest(target=target, options=options, findings=risk.findings, risk_summary=risk.risk_summary))
+    cfg = await agent.audit_configuration(ctx, AuditConfigurationRequest(target=target, options=options, inventory=inv["inventory"]))
+    rt = await agent.audit_runtime(ctx, AuditRuntimeRequest(target=target, options=options, inventory=inv["inventory"]))
+    risk = await agent.assess_risk(ctx, AssessRiskRequest(target=target, options=options, configuration_findings=cfg["findings"], runtime_findings=rt["findings"]))
+    rem = await agent.generate_remediation(ctx, GenerateRemediationRequest(target=target, options=options, findings=risk["findings"], risk_summary=risk["risk_summary"]))
 
-    assert inv.mutations_performed is False
-    assert cfg.mutations_performed is False
-    assert rt.mutations_performed is False
-    assert risk.mutations_performed is False
-    assert rem.mutations_performed is False
+    assert inv["mutations_performed"] is False
+    assert cfg["mutations_performed"] is False
+    assert rt["mutations_performed"] is False
+    assert risk["mutations_performed"] is False
+    assert rem["mutations_performed"] is False
     expected = {
         "outputs/compliance/clean/inventory.json",
         "outputs/compliance/clean/findings.json",
@@ -58,7 +58,7 @@ async def test_clean_synthetic_audit_outputs_expected_files() -> None:
         "outputs/compliance/clean/sarif.json",
     }
     assert expected.issubset(set(ctx.artifacts))
-    assert "sk-" not in rem.remediation_markdown
+    assert "sk-" not in rem["remediation_markdown"]
 
 
 @pytest.mark.asyncio
@@ -74,14 +74,14 @@ async def test_noncompliant_secret_and_public_exposure_findings_are_redacted() -
     options = AuditOptions(audit_id="bad", probe_mode="none", max_files=10)
 
     inv = await agent.inventory_agent(ctx, InventoryRequest(target=target, options=options))
-    cfg = await agent.audit_configuration(ctx, AuditConfigurationRequest(target=target, options=options, inventory=inv.inventory))
+    cfg = await agent.audit_configuration(ctx, AuditConfigurationRequest(target=target, options=options, inventory=inv["inventory"]))
 
-    controls = {f.control_id for f in cfg.findings}
+    controls = {f["control_id"] for f in cfg["findings"]}
     assert "A2A-SECRET-001" in controls
     assert "A2A-NET-001" in controls
     assert "A2A-UNTRUSTED-001" in controls
-    assert "sk-abcdefghijklmnopqrstuvwxyz" not in str(inv.inventory)
-    assert "REDACTED" in str(inv.inventory)
+    assert "sk-abcdefghijklmnopqrstuvwxyz" not in str(inv["inventory"])
+    assert "REDACTED" in str(inv["inventory"])
 
 
 @pytest.mark.asyncio
@@ -93,10 +93,10 @@ async def test_unreachable_or_disallowed_url_is_not_fetched_and_is_structured()
 
     inv = await agent.inventory_agent(ctx, InventoryRequest(target=target, options=options))
 
-    assert inv.status in {"ok", "partial"}
-    assert inv.inventory["agent_card"]["status"] == "not_tested"
-    assert inv.inventory["deployment_health"]["status"] == "not_tested"
-    assert any("allowlisted" in warning or "unsafe" in warning for warning in inv.warnings)
+    assert inv["status"] in {"ok", "partial"}
+    assert inv["inventory"]["agent_card"]["status"] == "not_tested"
+    assert inv["inventory"]["deployment_health"]["status"] == "not_tested"
+    assert any("allowlisted" in warning or "unsafe" in warning for warning in inv["warnings"])
 
 
 @pytest.mark.asyncio
@@ -113,10 +113,10 @@ async def test_malicious_repository_content_is_evidence_not_instruction() -> Non
     options = AuditOptions(audit_id="malicious", probe_mode="none", max_files=10)
 
     inv = await agent.inventory_agent(ctx, InventoryRequest(target=target, options=options))
-    cfg = await agent.audit_configuration(ctx, AuditConfigurationRequest(target=target, options=options, inventory=inv.inventory))
+    cfg = await agent.audit_configuration(ctx, AuditConfigurationRequest(target=target, options=options, inventory=inv["inventory"]))
 
-    assert inv.inventory["repository"]["suspicious_instruction_count"] == 1
-    assert "A2A-UNTRUSTED-001" in {f.control_id for f in cfg.findings}
+    assert inv["inventory"]["repository"]["suspicious_instruction_count"] == 1
+    assert "A2A-UNTRUSTED-001" in {f["control_id"] for f in cfg["findings"]}
     event_kinds = [event.kind for event in ctx.events]
     assert "audit_started" in event_kinds
     assert "audit_completed" in event_kinds