neuron-tai/tests/test_openai_gateway.py

"""US-005 integration tests: OpenAI-compatible gateway."""

import json
import urllib.error
import urllib.request

import pytest

from meshnet_gateway.server import GatewayServer
from meshnet_node.server import StubNodeServer
from meshnet_tracker.server import TrackerServer


def _register_node(tracker_url: str, endpoint: str, shard_start: int, shard_end: int) -> None:
    data = json.dumps({
        "endpoint": endpoint,
        "shard_start": shard_start,
        "shard_end": shard_end,
        "hardware_profile": {},
        "score": 1.0,
    }).encode()
    req = urllib.request.Request(
        f"{tracker_url}/v1/nodes/register",
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req) as r:
        r.read()


@pytest.fixture
def two_node_setup():
    """Start tracker, two stub nodes, and gateway; yield (gateway_url, node_a, node_b)."""
    tracker = TrackerServer()
    tracker_port = tracker.start()
    tracker_url = f"http://127.0.0.1:{tracker_port}"

    node_a = StubNodeServer(shard_start=0, shard_end=15, is_last_shard=False)
    port_a = node_a.start()

    node_b = StubNodeServer(shard_start=16, shard_end=31, is_last_shard=True)
    port_b = node_b.start()

    _register_node(tracker_url, f"http://127.0.0.1:{port_a}", 0, 15)
    _register_node(tracker_url, f"http://127.0.0.1:{port_b}", 16, 31)

    gateway = GatewayServer(tracker_url=tracker_url)
    gateway_port = gateway.start()
    gateway_url = f"http://127.0.0.1:{gateway_port}"

    yield gateway_url, node_a, node_b

    gateway.stop()
    node_a.stop()
    node_b.stop()
    tracker.stop()


def test_health_check(two_node_setup):
    """GET /v1/health returns 200."""
    gateway_url, _, _ = two_node_setup
    with urllib.request.urlopen(f"{gateway_url}/v1/health") as resp:
        assert resp.status == 200


def test_get_models_returns_preset_list(two_node_setup):
    """GET /v1/models returns OpenAI-format list containing stub-model."""
    gateway_url, _, _ = two_node_setup
    with urllib.request.urlopen(f"{gateway_url}/v1/models") as resp:
        body = json.loads(resp.read())
    assert body["object"] == "list"
    model_ids = [m["id"] for m in body["data"]]
    assert "stub-model" in model_ids


def test_non_streaming_via_openai_sdk(two_node_setup):
    """OpenAI SDK returns a valid ChatCompletion from the gateway (non-streaming)."""
    import openai
    gateway_url, _, _ = two_node_setup
    client = openai.OpenAI(base_url=f"{gateway_url}/v1", api_key="test")
    response = client.chat.completions.create(
        model="stub-model",
        messages=[{"role": "user", "content": "hello"}],
    )
    assert response.choices[0].message.content == "stub response to: hello"
    assert response.choices[0].finish_reason == "stop"


def test_streaming_via_openai_sdk(two_node_setup):
    """stream=True delivers text/event-stream chunks parseable by the OpenAI SDK."""
    import openai
    gateway_url, _, _ = two_node_setup
    client = openai.OpenAI(base_url=f"{gateway_url}/v1", api_key="test")

    chunks = []
    with client.chat.completions.create(
        model="stub-model",
        messages=[{"role": "user", "content": "hello"}],
        stream=True,
    ) as stream:
        for chunk in stream:
            if chunk.choices and chunk.choices[0].delta.content:
                chunks.append(chunk.choices[0].delta.content)

    assert "stub response to: hello" in "".join(chunks)


def test_unavailable_model_returns_openai_format_503(two_node_setup):
    """Unknown model → HTTP 503 with OpenAI-format error body (code='model_not_available')."""
    gateway_url, _, _ = two_node_setup
    payload = json.dumps({
        "model": "nonexistent-model",
        "messages": [{"role": "user", "content": "hello"}],
    }).encode()
    req = urllib.request.Request(
        f"{gateway_url}/v1/chat/completions",
        data=payload,
        headers={"Content-Type": "application/json", "Authorization": "Bearer test"},
        method="POST",
    )
    try:
        urllib.request.urlopen(req)
        raise AssertionError("Expected 503 for unknown model")
    except urllib.error.HTTPError as exc:
        assert exc.code == 503
        body = json.loads(exc.read())
        assert isinstance(body.get("error"), dict), "Expected OpenAI-format error dict"
        assert body["error"]["code"] == "model_not_available"


def test_langchain_chat_openai(two_node_setup):
    """LangChain ChatOpenAI works against the gateway."""
    from langchain_openai import ChatOpenAI
    gateway_url, _, _ = two_node_setup
    llm = ChatOpenAI(
        base_url=f"{gateway_url}/v1",
        api_key="test",
        model="stub-model",
    )
    result = llm.invoke("hello")
    assert "stub response to: hello" in result.content


def test_streaming_end_to_end_http(two_node_setup):
    """End-to-end streaming at HTTP level: SSE format, DONE sentinel, content chunk."""
    gateway_url, _, node_b = two_node_setup
    payload = json.dumps({
        "model": "stub-model",
        "messages": [{"role": "user", "content": "stream test"}],
        "stream": True,
    }).encode()
    req = urllib.request.Request(
        f"{gateway_url}/v1/chat/completions",
        data=payload,
        headers={"Content-Type": "application/json", "Authorization": "Bearer test"},
        method="POST",
    )
    with urllib.request.urlopen(req) as resp:
        assert resp.status == 200
        content_type = resp.headers["Content-Type"]
        assert "text/event-stream" in content_type
        raw = resp.read().decode()

    data_lines = [l for l in raw.strip().splitlines() if l.startswith("data: ")]
    assert data_lines, "No SSE data lines found"
    assert data_lines[-1] == "data: [DONE]"

    content_found = any(
        json.loads(line[6:]).get("choices", [{}])[0].get("delta", {}).get("content")
        for line in data_lines[:-1]
        if line != "data: [DONE]"
    )
    assert content_found, "No content delta found in SSE stream"
    assert node_b.received_activations, "node B did not participate in the pipeline"