TieMeasureFlow/src/backend/tests/test_vision_station_worker.py

"""#45: il worker di visione si sceglie per stazione.

La stazione si legge dalla sessione di produzione (`production_run_id`): con un
URL sulla stazione `_call_worker` lo riceve e il VisionResult lo registra;
senza URL o senza run si resta sul worker centrale, senza ripiego fra i due.
"""
import io

import httpx
import pytest
from fastapi import HTTPException
from PIL import Image

from src.backend.config import settings
from src.backend.models.api.production import ProductionRunCreate
from src.backend.models.orm.station import Station
from src.backend.models.orm.vision import VisionResult
from src.backend.services import production_service
from src.backend.services.vision_service import _call_worker
from src.backend.tests.conftest import auth_headers, create_test_recipe
from src.backend.tests.test_vision_execute import recipe_parts
from sqlalchemy import select

STATION_URL = "http://10.0.0.7:8100"


def an_image() -> bytes:
    buffer = io.BytesIO()
    Image.new("L", (64, 64), color=0).save(buffer, format="PNG")
    return buffer.getvalue()


@pytest.fixture
def worker_calls(monkeypatch):
    """Fake `_call_worker` che ricorda i kwargs con cui e' stato chiamato."""
    calls: list[dict] = []

    async def fake_call(image_bytes: bytes, graph: dict, **kwargs) -> dict:
        calls.append(kwargs)
        return {"outputs": {"t1.diameter": 10.0}, "failures": [],
                "engine_version": "e" * 40, "duration_ms": 3}

    monkeypatch.setattr("src.backend.services.vision_service._call_worker", fake_call)
    return calls


async def _camera_task(db_session, user_id: int):
    recipe = await create_test_recipe(db_session, user_id)
    version, task, subtask = await recipe_parts(db_session, recipe.id)
    task.task_type = "camera_measure"
    task.vision_json = {"schema_version": 1, "name": "p", "tools": []}
    subtask.vision_output = "t1.diameter"
    await db_session.commit()
    return recipe, task


async def _station_with_run(db_session, user, recipe, worker_url):
    station = Station(
        code="ST-W", name="Stazione W", active=True, created_by=user.id,
        vision_worker_url=worker_url,
    )
    db_session.add(station)
    await db_session.commit()
    await db_session.refresh(station)
    run = await production_service.open_run(
        db_session,
        ProductionRunCreate(station_code=station.code, recipe_id=recipe.id),
        user,
    )
    await db_session.commit()
    return station, run


async def _execute(client, user, task, run=None):
    data = {"task_id": str(task.id)}
    if run is not None:
        data["production_run_id"] = str(run.id)
    response = await client.post(
        "/api/vision/execute",
        headers=auth_headers(user),
        files={"image": ("p.png", an_image(), "image/png")},
        data=data,
    )
    assert response.status_code == 200, response.text


async def _the_result(db_session) -> VisionResult:
    return (await db_session.execute(select(VisionResult))).scalars().one()


@pytest.mark.asyncio
async def test_station_with_url_gets_the_call_and_the_record(
    client, db_session, measurement_tec_user, worker_calls,
):
    recipe, task = await _camera_task(db_session, measurement_tec_user.id)
    station, run = await _station_with_run(
        db_session, measurement_tec_user, recipe, STATION_URL)

    await _execute(client, measurement_tec_user, task, run)

    assert worker_calls[0]["worker_url"] == STATION_URL
    result = await _the_result(db_session)
    assert result.executed_on == "station"
    assert result.worker_url == STATION_URL
    assert result.station_id == station.id


@pytest.mark.asyncio
async def test_station_without_url_uses_the_central_worker(
    client, db_session, measurement_tec_user, worker_calls,
):
    recipe, task = await _camera_task(db_session, measurement_tec_user.id)
    station, run = await _station_with_run(
        db_session, measurement_tec_user, recipe, None)

    await _execute(client, measurement_tec_user, task, run)

    assert worker_calls[0]["worker_url"] is None
    result = await _the_result(db_session)
    assert result.executed_on == "server"
    assert result.worker_url == settings.vision_worker_url.rstrip("/")
    assert result.station_id == station.id


@pytest.mark.asyncio
async def test_without_a_run_the_central_worker_is_used(
    client, db_session, measurement_tec_user, worker_calls,
):
    _, task = await _camera_task(db_session, measurement_tec_user.id)

    await _execute(client, measurement_tec_user, task)

    assert worker_calls[0]["worker_url"] is None
    result = await _the_result(db_session)
    assert result.executed_on == "server"
    assert result.worker_url == settings.vision_worker_url.rstrip("/")
    assert result.station_id is None


class _Response:
    status_code = 200

    def json(self):
        return {"outputs": {}, "failures": [], "engine_version": "e" * 40}


class _RecordingClient:
    """httpx.AsyncClient finto: registra gli URL postati, opzionalmente fallisce."""

    def __init__(self, posts: list, raises=None):
        self._posts = posts
        self._raises = raises

    async def __aenter__(self):
        return self

    async def __aexit__(self, *exc_info):
        return False

    async def post(self, url, *args, **kwargs):
        self._posts.append(url)
        if self._raises is not None:
            raise self._raises
        return _Response()


def _install(monkeypatch, raises=None) -> list:
    posts: list = []
    monkeypatch.setattr(
        "src.backend.services.vision_service.httpx.AsyncClient",
        lambda *a, **k: _RecordingClient(posts, raises),
    )
    return posts


@pytest.mark.asyncio
async def test_call_worker_posts_to_the_given_url(monkeypatch):
    posts = _install(monkeypatch)
    # lo "/" finale non deve produrre "//run"
    await _call_worker(b"img", {"tools": []}, worker_url=STATION_URL + "/")
    assert posts == [f"{STATION_URL}/run"]


@pytest.mark.asyncio
async def test_call_worker_defaults_to_the_central_url(monkeypatch):
    posts = _install(monkeypatch)
    await _call_worker(b"img", {"tools": []})
    assert posts == [f"{settings.vision_worker_url.rstrip('/')}/run"]


@pytest.mark.asyncio
async def test_station_worker_down_is_502_with_its_url_and_no_fallback(monkeypatch):
    posts = _install(monkeypatch, raises=httpx.ConnectError("refused"))
    with pytest.raises(HTTPException) as exc:
        await _call_worker(b"img", {"tools": []}, worker_url=STATION_URL)
    assert exc.value.status_code == 502
    assert STATION_URL in exc.value.detail
    # una sola POST: il centrale non viene tentato
    assert posts == [f"{STATION_URL}/run"]