atlas-iac/testing/tests/test_hermes_model_gate_lan.py

"""Exercise the LAN gateway's authentication and local-only HTTP boundary."""

import http.client
import io
import json
from pathlib import Path
import threading
from urllib.error import URLError

import pytest
import yaml
from services.hermes.scripts import lan_generate
import sys

sys.modules.setdefault("lan_generate", lan_generate)


@pytest.fixture
def gateway(tmp_path, monkeypatch):
    """Run the shipped handler with an isolated token and intercepted upstream."""
    path = Path(__file__).parents[2] / "services/hermes/model-gate-configmap.yaml"
    namespace = {"__name__": "test_gateway"}
    exec(compile(yaml.safe_load(path.read_text())["data"]["model_gate.py"], str(path), "exec"), namespace)
    token = tmp_path / "token"
    token.write_text("a" * 64)
    namespace["LAN_API_TOKEN_FILE"] = token
    calls = []

    def upstream(path, payload=None, timeout=10):
        calls.append((path, payload, timeout))
        if path == "/api/version":
            return {"version": lan_generate.RUNTIME}
        if path == "/api/tags":
            return {"models": [{"name": lan_generate.MODEL, "digest": lan_generate.DIGEST}]}
        return {"model": lan_generate.MODEL, "response": "LAN_OK", "done": True,
                "context": [1, 2, 3], "thinking": "never forward this"}

    monkeypatch.setattr(lan_generate, "_request", upstream)
    monkeypatch.setattr(lan_generate, "_inference", threading.BoundedSemaphore(1))
    server = namespace["ThreadingHTTPServer"](("127.0.0.1", 0), namespace["LanHandler"])
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()

    def request(method="GET", path="/healthz", body=None, **headers):
        defaults = {"Authorization": "Bearer " + "a" * 64, "X-Forwarded-For": "192.168.22.8"}
        defaults.update(headers)
        connection = http.client.HTTPConnection(*server.server_address, timeout=3)
        connection.request(method, path, body=body, headers=defaults)
        response = connection.getresponse()
        result = response.status, json.loads(response.read())
        connection.close()
        return result

    yield namespace, calls, request
    server.shutdown()
    server.server_close()
    thread.join(timeout=2)


@pytest.mark.parametrize("address", ["", "203.0.113.1", "192.168.22.evil", "192.168.22.8, 203.0.113.1"])
def test_non_lan_and_spoofed_addresses_are_rejected(gateway, address):
    _, calls, request = gateway
    assert request(**{"X-Forwarded-For": address})[0] == 403
    assert not calls


@pytest.mark.parametrize("token", ["", "Bearer wrong", "Basic abc", "Bearer é"])
def test_token_is_required(gateway, token):
    _, calls, request = gateway
    assert request(Authorization=token)[0] == 401
    assert not calls


def test_health_requires_a_readable_credential(gateway):
    namespace, calls, request = gateway
    assert request()[0] == 200
    namespace["LAN_API_TOKEN_FILE"].unlink()
    assert request()[0] == 503
    assert len(calls) == 2


def test_generation_uses_only_local_upstream_without_forwarding_token(gateway):
    namespace, calls, request = gateway
    body = json.dumps({"model": namespace["LAN_MODEL"], "prompt": "Say LAN_OK", "stream": False})
    status, result = request("POST", "/api/generate", body)
    assert status == 200 and result["response"] == "LAN_OK"
    assert "context" not in result and "thinking" not in result
    assert result["inference_provenance"]["model_digest"] == lan_generate.DIGEST
    assert [call[0] for call in calls] == ["/api/version", "/api/tags", "/api/generate"]
    assert calls[-1][1]["options"]["num_predict"] == 256
    assert calls[-1][1]["options"]["num_ctx"] == 8192
    assert 1190 < calls[-1][2] <= 1200


@pytest.mark.parametrize("changes", [
    {"model": "external/model"}, {"stream": True}, {"tools": []}, {"images": []},
    {"context": []}, {"keep_alive": 0}, {"options": {"num_predict": -1}},
    {"options": {"num_predict": 2049}}, {"options": {"num_gpu": 100}},
    {"options": {"temperature": float("nan")}}, {"prompt": []},
])
def test_unsupported_payloads_never_reach_ollama(gateway, changes):
    namespace, calls, request = gateway
    payload = {"model": namespace["LAN_MODEL"], "prompt": "hello", "stream": False, **changes}
    assert request("POST", "/api/generate", json.dumps(payload))[0] == 400
    assert not calls


@pytest.mark.parametrize("length,status", [("-1", 413), ("131073", 413), ("invalid", 400)])
def test_invalid_request_lengths_are_rejected_before_read(gateway, length, status):
    _, calls, request = gateway
    assert request("POST", "/api/generate", "{}", **{"Content-Length": length})[0] == status
    assert not calls


def test_outage_and_concurrent_request_fail_without_fallback(gateway, monkeypatch):
    namespace, calls, request = gateway
    body = json.dumps({"model": namespace["LAN_MODEL"], "prompt": "hello", "stream": False})
    lan_generate._inference.acquire()
    assert request("POST", "/api/generate", body)[0] == 429
    lan_generate._inference.release()

    def unavailable(*args, **kwargs):
        raise URLError("local service unavailable")

    monkeypatch.setattr(lan_generate, "_request", unavailable)
    assert request("POST", "/api/generate", body)[0] == 503
    assert lan_generate._inference.acquire(timeout=1)
    lan_generate._inference.release()
    assert not calls


@pytest.mark.parametrize("path", ["/api/pull", "/api/delete", "/v1/chat/completions", "/api/chat"])
def test_model_management_and_agent_routes_are_not_exposed(gateway, path):
    _, calls, request = gateway
    assert request("POST", path, "{}")[0] == 404
    assert not calls


@pytest.mark.parametrize("path", ["/api/batch/models", "/api/batch/generate", "/terminal", "/api/tags"])
def test_only_initial_pilot_operations_are_exposed(gateway, path):
    _, calls, request = gateway
    assert request(path=path)[0] == 404
    assert request("POST", path, "{}")[0] == 404
    assert not calls


def test_logs_and_errors_never_echo_request_material(gateway, capsys):
    namespace, calls, request = gateway
    marker = "PRIVATE_CANARY_TEST_ONLY"
    assert request(path="/unknown?prompt=" + marker)[0] == 404
    body = json.dumps({"model": namespace["LAN_MODEL"], "prompt": marker,
                       "stream": False, marker: "unknown field"})
    status, error = request("POST", "/api/generate", body)
    assert status == 400 and marker not in json.dumps(error)
    assert marker not in capsys.readouterr().out
    assert not calls


def test_context_admission_does_not_truncate_or_call_backend(gateway):
    namespace, calls, request = gateway
    body = json.dumps({"model": namespace["LAN_MODEL"], "prompt": "x" * 8000, "stream": False})
    assert request("POST", "/api/generate", body)[0] == 413
    assert not calls


def test_structured_output_and_exact_model_pins(gateway, monkeypatch):
    _, calls, request = gateway
    original = lan_generate._request
    schema = {"type": "object", "properties": {"status": {"type": "string"}}, "required": ["status"]}
    body = json.dumps({"model": lan_generate.MODEL, "prompt": "Return JSON status ok",
                       "stream": False, "format": schema})

    def structured(path, payload=None, timeout=10):
        result = original(path, payload, timeout)
        if path == "/api/generate":
            result["response"] = '{"status":"ok"}'
        return result

    monkeypatch.setattr(lan_generate, "_request", structured)
    status, result = request("POST", "/api/generate", body)
    assert status == 200 and json.loads(result["response"]) == {"status": "ok"}
    assert calls[-1][1]["format"] == schema
    calls.clear()
    monkeypatch.setattr(lan_generate, "RUNTIME", "mismatch")
    # The fixture follows RUNTIME, so force an unexpected version independently.
    monkeypatch.setattr(lan_generate, "_request", lambda *a, **k: {"version": "different"})
    assert request("POST", "/api/generate", body)[0] == 503
    assert not calls