"""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]]}] CALIBRATION = [{"command": "make it toasty", "calls": [["living_room_heating", 24]]}] 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, "calibration": CALIBRATION, "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["calibration"], installation["tests"], ) == (EXAMPLES, CALIBRATION, 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"