Agent Advisor

scripts/test_verify_model_path.py

import json
import pathlib

import jsonschema
import pytest

import model_recommendation
import verify_model_path


NOW = "2026-07-20T12:00:00+00:00"


def _input(requirements=None):
    return {
        "schema_version": 2,
        "region": "us-east-1",
        "primary_unit": "support-agent",
        "workloads": [
            {
                "workload_id": "support-agent",
                "source": {
                    "provider": "anthropic",
                    "model_ids": ["claude-3-7-sonnet-latest"],
                    "sdk": "anthropic",
                    "api_surface": "messages",
                    "source_paths": ["src/agent.py"],
                },
                "requirements": {
                    "priority": "balanced",
                    "critical_features": ["tool_use"],
                    **(requirements or {}),
                },
                "detected_features": [],
            }
        ],
    }


class FakeMessages:
    def __init__(self, response=None, error=None):
        self.calls = []
        self.response = response or type(
            "Response", (), {"model": "anthropic.claude-sonnet-5"}
        )()
        self.error = error

    def create(self, **kwargs):
        self.calls.append(kwargs)
        if self.error:
            raise self.error
        return self.response


class FakeMantleClient:
    def __init__(self, messages=None):
        self.messages = messages or FakeMessages()


class FakeRuntimeClient:
    def __init__(self):
        self.converse_calls = []
        self.invoke_calls = []

    def converse(self, **kwargs):
        self.converse_calls.append(kwargs)
        return {"modelId": kwargs["modelId"]}

    def invoke_model(self, **kwargs):
        self.invoke_calls.append(kwargs)
        return {}


def _verify(requirements=None, **kwargs):
    recommendation = model_recommendation.recommend(_input(requirements))
    return verify_model_path.verify_recommendation(
        recommendation, now=NOW, **kwargs
    )


def test_mantle_probe_uses_recommended_clean_id():
    client = FakeMantleClient()
    result = _verify(mantle_client_factory=lambda region: client)
    verification = result["workloads"]["support-agent"]

    assert verification["status"] == "passed"
    assert verification["response_model_id"] == "anthropic.claude-sonnet-5"
    assert client.messages.calls[0]["model"] == "anthropic.claude-sonnet-5"
    assert client.messages.calls[0]["max_tokens"] == 8


def test_converse_probe_uses_resolved_cris_without_substitution():
    client = FakeRuntimeClient()
    result = _verify(
        {
            "governance": ["guardrails"],
            "data_residency": "global_allowed",
        },
        runtime_client_factory=lambda region: client,
    )
    verification = result["workloads"]["support-agent"]

    assert verification["status"] == "passed"
    assert (
        client.converse_calls[0]["modelId"]
        == "global.anthropic.claude-sonnet-5"
    )


def test_invoke_probe_uses_native_bedrock_body():
    client = FakeRuntimeClient()
    result = _verify(
        {
            "requires_native_payload": True,
            "data_residency": "global_allowed",
        },
        runtime_client_factory=lambda region: client,
    )
    verification = result["workloads"]["support-agent"]
    body = json.loads(client.invoke_calls[0]["body"])

    assert verification["status"] == "passed"
    assert body["anthropic_version"] == "bedrock-2023-05-31"
    assert body["messages"][0]["content"] == verify_model_path.PROMPT


def test_unresolved_runtime_cris_does_not_create_client():
    calls = []
    result = _verify(
        {"governance": ["guardrails"]},
        runtime_client_factory=lambda region: calls.append(region),
    )
    verification = result["workloads"]["support-agent"]

    assert verification["status"] == "needs_resolution"
    assert verification["invocation_model_id"] is None
    assert calls == []


def test_probe_failure_is_structured():
    client = FakeMantleClient(FakeMessages(error=PermissionError("denied")))
    result = _verify(mantle_client_factory=lambda region: client)
    verification = result["workloads"]["support-agent"]

    assert verification["status"] == "failed"
    assert verification["error"] == {
        "type": "PermissionError",
        "message": "denied",
    }


def test_decision_required_is_not_probed():
    result = _verify(
        {
            "preserve_messages_api": True,
            "governance": ["guardrails"],
        }
    )

    assert result["workloads"]["support-agent"]["status"] == "not_applicable"


def test_verification_output_matches_schema():
    result = _verify(mantle_client_factory=lambda region: FakeMantleClient())
    schema = json.loads(
        (
            pathlib.Path(verify_model_path.__file__).parent
            / "schemas"
            / "model-verification.json"
        ).read_text()
    )

    jsonschema.validate(
        result,
        schema,
        format_checker=jsonschema.FormatChecker(),
    )


# --- OpenAI provider verification (mocked, no network) ---------------------


class FakeResponses:
    def __init__(self, response=None, error=None):
        self.calls = []
        self.response = response or type("Resp", (), {"model": "openai.gpt-5.6-sol"})()
        self.error = error

    def create(self, **kwargs):
        self.calls.append(kwargs)
        if self.error:
            raise self.error
        return self.response


class FakeCompletions:
    def __init__(self, response=None, error=None):
        self.calls = []
        self.response = response or type("Resp", (), {"model": "openai.gpt-6-sol"})()
        self.error = error

    def create(self, **kwargs):
        self.calls.append(kwargs)
        if self.error:
            raise self.error
        return self.response


class FakeChat:
    def __init__(self, completions=None):
        self.completions = completions or FakeCompletions()


class FakeOpenAIClient:
    def __init__(self, responses=None, chat=None):
        self.responses = responses or FakeResponses()
        self.chat = chat or FakeChat()


@pytest.mark.parametrize("region", ["us-east-1", "us-west-2"])
@pytest.mark.parametrize("surface", ["chat_completions", "responses"])
def test_astra_recommendation_probes_selected_api_in_supported_regions(region, surface):
    data = _openai_input()
    data["region"] = region
    data["workloads"][0]["source"].update(
        model_ids=["gpt-6-astra"], api_surface=surface
    )
    recommendation = model_recommendation.recommend(data)
    completions = FakeResponses(response={"model": "openai.gpt-6-astra"})
    client = type("ChatClient", (), {
        "chat": type("Chat", (), {"completions": completions})(),
        "responses": completions,
    })()
    result = verify_model_path.verify_recommendation(
        recommendation, now=NOW,
        openai_responses_client_factory=lambda region: client,
        openai_chat_client_factory=lambda region: client,
    )
    verification = result["workloads"]["openai-svc"]
    assert verification["status"] == "passed"
    assert verification["api_path"] == (
        "mantle_openai_chat" if surface == "chat_completions" else "mantle_openai_responses"
    )
    assert verification["region"] == region
    assert completions.calls[0]["model"] == "openai.gpt-6-astra"
    if surface == "chat_completions":
        assert completions.calls[0]["messages"][0]["role"] == "user"
    else:
        assert completions.calls[0]["input"] == verify_model_path.PROMPT


@pytest.mark.parametrize("prefix", ["us", "global"])
@pytest.mark.parametrize("as_arn", [False, True])
def test_astra_system_profile_representation_reaches_verifier_unchanged(prefix, as_arn):
    model_id = f"{prefix}.openai.gpt-6-astra"
    if as_arn:
        model_id = f"arn:aws:bedrock:us-west-2:111122223333:inference-profile/{model_id}"
    data = _openai_input({"governance": ["guardrails"], "inference_profile_id": model_id})
    data["region"] = "us-west-2"
    data["workloads"][0]["source"]["model_ids"] = ["gpt-6-astra"]
    schemas = pathlib.Path(model_recommendation.__file__).parent / "schemas"
    jsonschema.validate(data, json.loads((schemas / "model-recommendation-input.json").read_text()))
    recommendation = model_recommendation.recommend(data)
    jsonschema.validate(recommendation, json.loads((schemas / "model-recommendation.json").read_text()))
    selected = recommendation["workloads"]["openai-svc"]
    assert selected["decision_status"] == "recommended"
    assert selected["invocation_model_id"] == model_id
    client = FakeRuntimeClient()
    result = verify_model_path.verify_recommendation(
        recommendation, now=NOW, runtime_client_factory=lambda region: client,
    )
    assert result["workloads"]["openai-svc"]["status"] == "passed"
    assert client.converse_calls[0]["modelId"] == model_id


def _openai_input(requirements=None, source=None, region="us-east-2"):
    source_data = {
        "provider": "openai",
        "model_ids": ["gpt-5.4"],
        "sdk": "openai",
        "api_surface": "responses",
        "source_paths": ["src/app.py"],
    }
    source_data.update(source or {})
    return {
        "schema_version": 2,
        "region": region,
        "primary_unit": "openai-svc",
        "workloads": [
            {
                "workload_id": "openai-svc",
                "source": source_data,
                "requirements": {
                    "priority": "balanced",
                    "critical_features": [],
                    **(requirements or {}),
                },
                "detected_features": [],
            }
        ],
    }


def _verify_openai(requirements=None, source=None, region="us-east-2", **kwargs):
    recommendation = model_recommendation.recommend(
        _openai_input(requirements, source, region)
    )
    return verify_model_path.verify_recommendation(recommendation, now=NOW, **kwargs)


OPUS55 = "anthropic.claude-opus-5-5"
PROFILE_ARN = f"arn:aws:bedrock:us-east-1:123456789012:inference-profile/us.{OPUS55}"
APPLICATION_ARN = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/test-profile"


def _opus_recommendation(provider, path, profile):
    data = _input() if provider == "anthropic" else _openai_input()
    data["region"] = "us-east-1"
    requirements = data["workloads"][0]["requirements"]
    requirements.update(priority="quality", data_residency="geo_required",
                        cris_geography="us", inference_profile_id=profile)
    if provider == "anthropic":
        requirements["preferred_api_path"] = path
    else:
        requirements.update(governance=["guardrails"], preserve_openai_api=False)
    recommendation = model_recommendation.recommend(data)
    jsonschema.validate(recommendation, json.loads(
        (pathlib.Path(verify_model_path.__file__).parent / "schemas/model-recommendation.json").read_text()
    ))
    return recommendation


@pytest.mark.parametrize("provider,path", [
    ("anthropic", "runtime_converse"), ("anthropic", "runtime_invoke"),
    ("openai", "runtime_converse"),
])
@pytest.mark.parametrize("profile", [f"us.{OPUS55}", PROFILE_ARN])
def test_system_profile_id_and_arn_reach_the_selected_probe(provider, path, profile):
    client = FakeRuntimeClient()
    result = verify_model_path.verify_recommendation(
        _opus_recommendation(provider, path, profile),
        runtime_client_factory=lambda region: client,
    )
    assert next(iter(result["workloads"].values()))["status"] == "passed"
    calls = client.converse_calls if path == "runtime_converse" else client.invoke_calls
    assert calls[0]["modelId"] == profile


@pytest.mark.parametrize("profile", [
    PROFILE_ARN.replace("us-east-1", "eu-west-1"),
    PROFILE_ARN.replace("arn:aws:", "arn:aws-us-gov:"),
    PROFILE_ARN.replace("us.anthropic", "global.anthropic"),
    PROFILE_ARN.replace("opus-5-5", "sonnet-5"),
])
def test_system_profile_arn_mismatches_do_not_reach_inference(profile):
    client = FakeRuntimeClient()
    result = verify_model_path.verify_recommendation(
        _opus_recommendation("anthropic", "runtime_converse", profile),
        runtime_client_factory=lambda region: client,
    )
    assert next(iter(result["workloads"].values()))["status"] == "needs_resolution"
    assert not client.converse_calls


def _profile_metadata(regions=("us-east-1", "us-west-2"), model=OPUS55, status="ACTIVE"):
    return {"status": status, "models": [
        {"modelArn": f"arn:aws:bedrock:{region}::foundation-model/{model}"}
        for region in regions
    ]}


@pytest.mark.parametrize("provider,path", [
    ("anthropic", "runtime_converse"), ("anthropic", "runtime_invoke"),
    ("openai", "runtime_converse"),
])
def test_application_profile_metadata_is_verified_before_invoking_its_arn(provider, path):
    client = FakeRuntimeClient()
    reads = []

    class Control:
        def get_inference_profile(self, inferenceProfileIdentifier):
            assert not client.converse_calls and not client.invoke_calls
            reads.append(inferenceProfileIdentifier)
            return _profile_metadata()

    recommendation = _opus_recommendation(provider, path, APPLICATION_ARN)
    assert next(iter(recommendation["workloads"].values()))["invocation_model_id"] == APPLICATION_ARN
    result = verify_model_path.verify_recommendation(
        recommendation, runtime_client_factory=lambda region: client,
        control_client_factory=lambda region: Control(),
    )
    assert next(iter(result["workloads"].values()))["status"] == "passed"
    assert reads == [APPLICATION_ARN, f"us.{OPUS55}"]
    calls = client.converse_calls if path == "runtime_converse" else client.invoke_calls
    assert calls[0]["modelId"] == APPLICATION_ARN


@pytest.mark.parametrize("metadata", [
    _profile_metadata(model="anthropic.claude-sonnet-5"),
    _profile_metadata(regions=("us-east-1", "eu-west-1")),
    _profile_metadata(regions=()),
    _profile_metadata(status="INACTIVE"),
])
def test_application_profile_wrong_model_or_destinations_fail_before_inference(metadata):
    client = FakeRuntimeClient()

    class Control:
        def get_inference_profile(self, inferenceProfileIdentifier):
            return metadata if inferenceProfileIdentifier == APPLICATION_ARN else _profile_metadata()

    result = verify_model_path.verify_recommendation(
        _opus_recommendation("anthropic", "runtime_converse", APPLICATION_ARN),
        runtime_client_factory=lambda region: client,
        control_client_factory=lambda region: Control(),
    )
    assert next(iter(result["workloads"].values()))["status"] == "failed"
    assert not client.converse_calls


def test_application_profile_metadata_denial_does_not_create_runtime_client():
    calls = []

    class Control:
        def get_inference_profile(self, **kwargs):
            raise PermissionError("GetInferenceProfile denied")

    result = verify_model_path.verify_recommendation(
        _opus_recommendation("anthropic", "runtime_converse", APPLICATION_ARN),
        runtime_client_factory=lambda region: calls.append(region),
        control_client_factory=lambda region: Control(),
    )
    verification = next(iter(result["workloads"].values()))
    assert verification["status"] == "failed"
    assert verification["error"]["type"] == "PermissionError"
    assert not calls


def test_application_profile_without_constraints_is_not_invoked():
    recommendation = _opus_recommendation("anthropic", "runtime_converse", APPLICATION_ARN)
    next(iter(recommendation["workloads"].values()))["verification"].pop("allowed_inference_profiles")
    calls = []
    result = verify_model_path.verify_recommendation(
        recommendation, runtime_client_factory=lambda region: calls.append(region),
        control_client_factory=lambda region: calls.append(region),
    )
    assert next(iter(result["workloads"].values()))["status"] == "failed"
    assert not calls


def test_global_application_profile_can_use_an_allowed_reference_after_an_unreadable_one():
    recommendation = model_recommendation.recommend(_input({
        "priority": "quality", "governance": ["guardrails"],
        "data_residency": "global_allowed", "inference_profile_id": APPLICATION_ARN,
    }))
    client = FakeRuntimeClient()

    class Control:
        def get_inference_profile(self, inferenceProfileIdentifier):
            if inferenceProfileIdentifier == f"us.{OPUS55}":
                raise PermissionError("US reference not readable")
            return _profile_metadata(regions=("us-east-1", "eu-west-1"))

    result = verify_model_path.verify_recommendation(
        recommendation, runtime_client_factory=lambda region: client,
        control_client_factory=lambda region: Control(),
    )
    assert next(iter(result["workloads"].values()))["status"] == "passed"
    assert client.converse_calls[0]["modelId"] == APPLICATION_ARN


def test_recommendation_cli_invalidates_passed_profile_cache_when_residency_changes(tmp_path):
    data = _input({
        "priority": "quality", "governance": ["guardrails"],
        "data_residency": "global_allowed", "inference_profile_id": APPLICATION_ARN,
    })
    source = tmp_path / "input.json"
    source.write_text(json.dumps(data))
    assert model_recommendation.main([str(source)]) == 0
    cache = tmp_path / "model-verification.json"
    passed = '{"workloads":{"support-agent":{"status":"passed"}}}\n'
    cache.write_text(passed)
    assert model_recommendation.main([str(source)]) == 0
    assert cache.read_text() == passed
    data["workloads"][0]["requirements"].update(data_residency="geo_required", cris_geography="us")
    source.write_text(json.dumps(data))
    assert model_recommendation.main([str(source)]) == 0
    assert not cache.exists()
    result = json.loads((tmp_path / "model-recommendation.json").read_text())
    workload = result["workloads"]["support-agent"]
    assert workload["invocation_model_id"] == APPLICATION_ARN
    assert workload["verification"]["allowed_inference_profiles"] == [f"us.{OPUS55}"]


def test_recommendation_cli_drops_orphaned_verification_before_first_write(tmp_path):
    source = tmp_path / "input.json"
    source.write_text(json.dumps(_input({"data_residency": "global_allowed"})))
    cache = tmp_path / "model-verification.json"
    cache.write_text('{"status":"passed"}')
    assert model_recommendation.main([str(source)]) == 0
    assert not cache.exists()


@pytest.mark.parametrize("provider,path", [
    ("anthropic", "runtime_converse"), ("anthropic", "runtime_invoke"),
    ("openai", "runtime_converse"),
])
@pytest.mark.parametrize("regions", [("eu-west-2",), ("eu-west-2", "eu-west-1")])
def test_london_application_profile_accepts_in_region_or_matching_geo_destinations(provider, path, regions):
    data = _input() if provider == "anthropic" else _openai_input()
    data["region"] = "eu-west-2"
    arn = APPLICATION_ARN.replace("us-east-1", "eu-west-2")
    requirements = data["workloads"][0]["requirements"]
    requirements.update(priority="quality", preferred_api_path=path, data_residency="geo_required",
                        cris_geography="eu", inference_profile_id=arn)
    if provider == "openai":
        requirements.update(governance=["guardrails"], preserve_openai_api=False)
    recommendation = model_recommendation.recommend(data)
    workload = next(iter(recommendation["workloads"].values()))
    assert workload["verification"]["allowed_in_region_model_arns"] == [
        f"arn:aws:bedrock:eu-west-2::foundation-model/{OPUS55}"
    ]
    client = FakeRuntimeClient()

    class Control:
        def get_inference_profile(self, inferenceProfileIdentifier):
            assert not client.converse_calls and not client.invoke_calls
            return _profile_metadata(regions=regions if inferenceProfileIdentifier == arn
                                     else ("eu-west-2", "eu-west-1"))

    result = verify_model_path.verify_recommendation(
        recommendation, runtime_client_factory=lambda region: client,
        control_client_factory=lambda region: Control(),
    )
    assert next(iter(result["workloads"].values()))["status"] == "passed"
    calls = client.invoke_calls if path == "runtime_invoke" else client.converse_calls
    assert calls[0]["modelId"] == arn


@pytest.mark.parametrize("metadata", [
    _profile_metadata(regions=("eu-west-1",)),
    _profile_metadata(regions=("eu-west-2",), model="anthropic.claude-sonnet-5"),
    _profile_metadata(regions=("eu-west-2",), status="INACTIVE"),
    _profile_metadata(regions=("eu-west-2", "us-east-1")),
])
def test_london_in_region_allowance_rejects_wrong_or_additional_destinations(metadata):
    arn = APPLICATION_ARN.replace("us-east-1", "eu-west-2")
    data = _input({"priority": "quality", "preferred_api_path": "runtime_converse",
                   "data_residency": "geo_required", "cris_geography": "eu",
                   "inference_profile_id": arn})
    data["region"] = "eu-west-2"
    client = FakeRuntimeClient()

    class Control:
        def get_inference_profile(self, inferenceProfileIdentifier):
            return metadata if inferenceProfileIdentifier == arn else _profile_metadata(
                regions=("eu-west-2", "eu-west-1"))

    result = verify_model_path.verify_recommendation(
        model_recommendation.recommend(data), runtime_client_factory=lambda region: client,
        control_client_factory=lambda region: Control(),
    )
    assert next(iter(result["workloads"].values()))["status"] == "failed"
    assert not client.converse_calls


def test_changed_in_region_allowance_invalidates_same_arn_verification(tmp_path):
    arn = APPLICATION_ARN.replace("us-east-1", "eu-west-2")
    data = _input({"priority": "quality", "preferred_api_path": "runtime_converse",
                   "data_residency": "geo_required", "cris_geography": "eu",
                   "inference_profile_id": arn})
    data["region"] = "eu-west-2"
    source = tmp_path / "input.json"
    source.write_text(json.dumps(data))
    assert model_recommendation.main([str(source)]) == 0
    cache = tmp_path / "model-verification.json"
    cache.write_text('{"status":"passed"}')
    catalog = model_recommendation.load_catalog()
    catalog["models"]["claude_opus_5_5"]["runtime_in_region_regions"] = []
    custom = tmp_path / "catalog.json"
    custom.write_text(json.dumps(catalog))
    assert model_recommendation.main([str(source), "--catalog", str(custom)]) == 0
    assert not cache.exists()
    result = json.loads((tmp_path / "model-recommendation.json").read_text())
    workload = next(iter(result["workloads"].values()))
    assert workload["invocation_model_id"] == arn
    assert workload["verification"]["allowed_inference_profiles"] == [f"eu.{OPUS55}"]
    assert "allowed_in_region_model_arns" not in workload["verification"]


def test_mantle_responses_probe_calls_responses_create_with_exact_id():
    client = FakeOpenAIClient()
    result = _verify_openai(
        {"api_continuity": "required"},
        openai_responses_client_factory=lambda region: client,
        openai_chat_client_factory=lambda region: client,
    )
    verification = result["workloads"]["openai-svc"]

    assert verification["status"] == "passed"
    assert verification["api_path"] == "mantle_openai_responses"
    assert client.responses.calls[0]["model"] == "openai.gpt-5.6-sol"
    assert client.responses.calls[0]["input"]
    assert verification["response_model_id"] == "openai.gpt-5.6-sol"


def test_mantle_responses_failure_does_not_substitute_model():
    client = FakeOpenAIClient(FakeResponses(error=RuntimeError("access denied")))
    result = _verify_openai(
        {"api_continuity": "required"},
        openai_responses_client_factory=lambda region: client,
        openai_chat_client_factory=lambda region: client,
    )
    verification = result["workloads"]["openai-svc"]

    assert verification["status"] == "failed"
    assert verification["invocation_model_id"] == "openai.gpt-5.6-sol"
    # Only the exact selected model was ever probed.
    assert [c["model"] for c in client.responses.calls] == ["openai.gpt-5.6-sol"]


def test_mantle_chat_probe_calls_chat_completions_with_exact_id():
    client = FakeOpenAIClient()
    result = _verify_openai(
        {"api_continuity": "required"},
        source={"model_ids": ["gpt-6-sol"], "api_surface": "chat_completions"},
        region="us-east-1",
        openai_chat_client_factory=lambda region: client,
    )
    verification = result["workloads"]["openai-svc"]
    schema = json.loads(
        (
            pathlib.Path(verify_model_path.__file__).parent
            / "schemas"
            / "model-verification.json"
        ).read_text()
    )
    jsonschema.validate(result, schema, format_checker=jsonschema.FormatChecker())

    assert verification["status"] == "passed"
    assert verification["api_path"] == "mantle_openai_chat"
    assert verification["response_model_id"] == "openai.gpt-6-sol"
    assert client.chat.completions.calls == [
        {
            "model": "openai.gpt-6-sol",
            "max_completion_tokens": 8,
            "messages": [{"role": "user", "content": verify_model_path.PROMPT}],
        }
    ]


def test_mantle_chat_failure_does_not_substitute_model():
    completions = FakeCompletions(error=RuntimeError("access denied"))
    client = FakeOpenAIClient(chat=FakeChat(completions))
    result = _verify_openai(
        {"api_continuity": "required"},
        source={"model_ids": ["gpt-6-sol"], "api_surface": "chat_completions"},
        region="us-east-1",
        openai_chat_client_factory=lambda region: client,
    )
    verification = result["workloads"]["openai-svc"]

    assert verification["status"] == "failed"
    assert verification["invocation_model_id"] == "openai.gpt-6-sol"
    assert [call["model"] for call in completions.calls] == ["openai.gpt-6-sol"]


def test_openai_decision_required_is_not_probed():
    called = {"n": 0}

    def factory(region):
        called["n"] += 1
        return FakeOpenAIClient()

    result = _verify_openai(
        {"api_continuity": "required", "governance": ["guardrails"]},
        openai_responses_client_factory=factory,
    )
    verification = result["workloads"]["openai-svc"]

    assert verification["status"] == "not_applicable"
    assert called["n"] == 0  # no client constructed for a non-selected recommendation