ewoooc/tests/test_nemotron_runtime_candidate_service.py

import json
import os
import subprocess
import sys

import pytest
import requests


class _Response:
    def __init__(self, payload):
        self._payload = payload

    def raise_for_status(self):
        return None

    def json(self):
        return self._payload


def test_candidate_chain_is_gcp_first_with_exact_digest_111_fallback():
    from services import nemotron_runtime_candidate_service as service

    candidates = service.build_nemotron_runtime_candidates()

    assert [candidate.label for candidate in candidates] == [
        "gcp_primary",
        "gcp_secondary",
        "ollama_111_fallback",
    ]
    assert [candidate.model for candidate in candidates] == [
        "qwen3:14b",
        "qwen3:14b",
        "qwen3:8b",
    ]
    assert candidates[0].is_fallback is False
    assert candidates[2].is_fallback is True
    assert candidates[0].request_timeout_sec == 60
    assert candidates[2].request_timeout_sec == 45
    assert candidates[0].num_predict == 2048
    assert candidates[2].num_predict == 512
    assert candidates[0].num_ctx is None
    assert candidates[2].num_ctx == 4096


def test_model_identity_fails_closed_on_digest_drift():
    from services import nemotron_runtime_candidate_service as service

    candidate = service.build_nemotron_runtime_candidates()[0]
    identity = service.inspect_nemotron_model_identity(
        candidate,
        request_get=lambda *_args, **_kwargs: _Response({
            "models": [{
                "name": candidate.model,
                "model": candidate.model,
                "digest": "deadbeef",
                "details": {},
            }]
        }),
    )

    assert identity["ok"] is False
    assert identity["digest_matches"] is False
    assert identity["error"] == "model_digest_mismatch"


def test_model_identity_reports_network_failure_without_loading_model():
    from services import nemotron_runtime_candidate_service as service

    candidate = service.build_nemotron_runtime_candidates()[2]

    def fail(*_args, **_kwargs):
        raise requests.Timeout("offline")

    identity = service.inspect_nemotron_model_identity(
        candidate,
        request_get=fail,
    )

    assert identity["ok"] is False
    assert identity["digest"] is None
    assert identity["error"].startswith("Timeout: offline")


@pytest.mark.parametrize(
    ("payload", "error_fragment"),
    [
        (None, "malformed_tags_payload"),
        ({"models": {}}, "malformed_models_schema"),
        ({"models": [None]}, "malformed_models_schema"),
        ({"models": [{"name": "other", "details": None}]}, "malformed_model_details"),
    ],
)
def test_model_identity_fails_closed_on_malformed_tags_schema(
    payload, error_fragment
):
    from services import nemotron_runtime_candidate_service as service

    candidate = service.build_nemotron_runtime_candidates()[0]
    identity = service.inspect_nemotron_model_identity(
        candidate,
        request_get=lambda *_args, **_kwargs: _Response(payload),
    )

    assert identity["ok"] is False
    assert identity["digest_matches"] is False
    assert error_fragment in str(identity["error"])


def test_model_identity_rejects_malformed_entry_mixed_with_valid_target():
    from services import nemotron_runtime_candidate_service as service

    candidate = service.build_nemotron_runtime_candidates()[0]
    identity = service.inspect_nemotron_model_identity(
        candidate,
        request_get=lambda *_args, **_kwargs: _Response({
            "models": [
                None,
                {
                    "name": candidate.model,
                    "model": candidate.model,
                    "digest": candidate.expected_digest,
                    "details": {},
                },
            ]
        }),
    )

    assert identity["ok"] is False
    assert identity["digest_matches"] is False
    assert identity["error"] == "malformed_models_schema"


def test_candidate_attempt_timeout_honors_env_but_111_resources_stay_fixed():
    env = os.environ.copy()
    env.update({
        "NEMOTRON_111_ATTEMPT_TIMEOUT_SEC": "33",
        "OLLAMA_111_NUM_CTX": "2048",
        "OLLAMA_111_NUM_PREDICT": "256",
    })
    completed = subprocess.run(
        [
            sys.executable,
            "-c",
            (
                "import json; "
                "from services.nemotron_runtime_candidate_service import "
                "build_nemotron_runtime_candidates; "
                "c=build_nemotron_runtime_candidates()[-1]; "
                "print(json.dumps(c.as_public_dict(), sort_keys=True))"
            ),
        ],
        cwd=os.path.dirname(os.path.dirname(__file__)),
        env=env,
        capture_output=True,
        check=False,
        text=True,
    )

    assert completed.returncode == 0, completed.stderr
    fallback = json.loads(completed.stdout.strip())
    assert fallback["request_timeout_sec"] == 33
    assert fallback["num_ctx"] == 4096
    assert fallback["num_predict"] == 512


@pytest.mark.parametrize(
    ("num_ctx", "num_predict"),
    [
        ("99999", "99999"),
        ("-7", "-9"),
        ("not-a-number", "invalid"),
    ],
)
def test_candidate_resource_contract_cannot_be_weakened_by_environment(
    num_ctx, num_predict
):
    env = os.environ.copy()
    env.update({
        "OLLAMA_111_NUM_CTX": num_ctx,
        "OLLAMA_111_NUM_PREDICT": num_predict,
    })
    completed = subprocess.run(
        [
            sys.executable,
            "-c",
            (
                "import json; "
                "from services.nemotron_runtime_candidate_service import "
                "build_nemotron_runtime_candidates; "
                "c=build_nemotron_runtime_candidates()[-1]; "
                "print(json.dumps(c.as_public_dict(), sort_keys=True))"
            ),
        ],
        cwd=os.path.dirname(os.path.dirname(__file__)),
        env=env,
        capture_output=True,
        check=False,
        text=True,
    )

    assert completed.returncode == 0, completed.stderr
    fallback = json.loads(completed.stdout.strip())
    assert fallback["num_ctx"] == 4096
    assert fallback["num_predict"] == 512