app/backend/tests/api/routes/test_wren_train.py

"""Train on a wren node: its wiring and settings go to wren's training flow.

The model package is optional, so ``FakeWren`` (tests/flow/conftest.py)
stands in for it.
"""

import json
from collections.abc import Generator
from datetime import UTC, datetime, timedelta

import pytest
from fastapi.testclient import TestClient
from sqlmodel import Session

from fluksio.core.config import settings
from fluksio.core.db import engine as db_engine
from fluksio.models import Run, RunArtifact
from tests.flow.conftest import FakeWren

FLOW = "wren_house"
PREFIX = f"{settings.API_V1_STR}/flows/{FLOW}/nodes"
EXAMPLES = [{"command": "warmer please", "calls": [["living_room_heating", 22]]}]
TESTS = [
    {"command": "is it cold outside", "calls": [["outdoor_temperature"]]},
    {"command": "good night", "calls": []},
]


@pytest.fixture(autouse=True)
def fake_wren(monkeypatch: pytest.MonkeyPatch) -> FakeWren:
    fake = FakeWren()
    monkeypatch.setattr("fluksio.flow.nodes.wren._wren", lambda: fake)
    return fake


@pytest.fixture(scope="module", autouse=True)
def house(
    client: TestClient, superuser_token_headers: dict[str, str]
) -> Generator[None, None, None]:
    """Two wren nodes, so a run for one is not the other's, and a function."""
    prompt = {"name": "command", "dtype": "str", "port": "prompt"}
    nodes = [
        {
            "id": "control",
            "type": "wren",
            "params": {
                "targets": {
                    "living_room_heating": {"range": [5, 30], "unit": "°C"},
                    "outdoor_temperature": {"description": "temperature outside"},
                },
                "examples": EXAMPLES,
                "tests": TESTS,
            },
            "requires": [
                prompt,
                {"name": "outdoor", "dtype": "float", "port": "outdoor_temperature"},
            ],
            "provides": [
                {"name": "setpoint", "dtype": "float", "port": "living_room_heating"}
            ],
        },
        {
            "id": "lights",
            "type": "wren",
            "requires": [prompt],
            "provides": [{"name": "kitchen_light", "dtype": "bool"}],
        },
        {"id": "sensor", "type": "python"},
    ]
    url = f"{settings.API_V1_STR}/flows/{FLOW}"
    saved = client.put(
        url, headers=superuser_token_headers, json={"name": FLOW, "nodes": nodes}
    )
    assert saved.status_code == 200, saved.text
    yield
    client.delete(url, headers=superuser_token_headers)


def test_train_submits_the_node_as_wren_s_installation(
    client, superuser_token_headers, monkeypatch
):
    """The draft's ports, targets, examples and tests, never from the cache."""
    seen: dict[str, object] = {}

    class Recorder:
        def submit(self, name, **kwargs):
            seen.update(kwargs, flow=name)
            return Run(id="train-1", flow=name, created_at=datetime.now(UTC))

    monkeypatch.setattr(client.app.state, "run_service", Recorder())
    answer = client.post(f"{PREFIX}/control/train", headers=superuser_token_headers)

    assert answer.status_code == 202, answer.text
    assert answer.json()["id"] == "train-1"
    assert (seen["flow"], seen["no_cache"], seen["cause"], seen["for_node"]) == (
        "wren_train",
        True,
        "train",
        f"{FLOW}.control",
    )
    installation = json.loads(seen["params"]["installation"])  # type: ignore[index]
    assert installation["targets"] == [
        {
            "name": "living_room_heating",
            "type": "float",
            "values": None,
            "range": [5, 30],
            "description": "",
            "unit": "°C",
            "readonly": False,
            "value": 5.0,
        },
        {
            "name": "outdoor_temperature",
            "type": "float",
            "values": None,
            "range": None,
            "description": "temperature outside",
            "unit": None,
            "readonly": True,
            "value": 0.0,
        },
    ]
    assert (installation["examples"], installation["tests"]) == (EXAMPLES, TESTS)


def test_train_says_how_to_register_a_missing_training_flow(
    client, superuser_token_headers
):
    answer = client.post(f"{PREFIX}/control/train", headers=superuser_token_headers)

    assert answer.status_code == 409
    assert "fluksio sync wren.flow" in answer.json()["detail"]


def test_only_a_wren_node_trains(client, superuser_token_headers):
    answer = client.post(f"{PREFIX}/sensor/train", headers=superuser_token_headers)
    assert answer.status_code == 400
    missing = client.post(f"{PREFIX}/nobody/train", headers=superuser_token_headers)
    assert missing.status_code == 404


def test_a_node_reads_its_own_latest_training(
    client, superuser_token_headers, fake_wren
):
    """With the model's digest and what its record says of it."""
    pack = client.app.state.artifact_store.put([b"weights"], name="wren.bin")
    fake_wren.pack = {"meta": {"threshold": 0.93, "test_exact_match": 0.5}}
    now = datetime.now(UTC)
    with Session(db_engine) as session:
        for run_id, node, ago in (
            ("train-old", "control", 3),
            ("train-new", "control", 2),
            ("train-lights", "lights", 1),
        ):
            session.add(
                Run(
                    id=run_id,
                    flow="wren_train",
                    status="ok",
                    for_node=f"{FLOW}.{node}",
                    result={"exact_match": 0.98, "coverage": 0.9},
                    created_at=now - timedelta(minutes=ago),
                )
            )
        session.add(
            RunArtifact(
                run_id="train-new", name="wren_train.pack", digest=pack["digest"]
            )
        )
        session.commit()

    training = client.get(f"{PREFIX}/control/train", headers=superuser_token_headers)

    assert training.status_code == 200
    assert training.json() == {
        "id": "train-new",
        "status": "ok",
        "error": "",
        "result": {"exact_match": 0.98, "coverage": 0.9},
        "model": pack["digest"],
        "threshold": 0.93,
        "test_exact_match": 0.5,
    }
    lights = client.get(f"{PREFIX}/lights/train", headers=superuser_token_headers)
    assert lights.json()["id"] == "train-lights"