from datetime import datetime from unittest.mock import AsyncMock, patch from fastapi.testclient import TestClient from open_notebook.domain.transformation import Transformation def _client() -> TestClient: from api.main import app return TestClient(app) def _transformation(model_id: str | None = None) -> Transformation: return Transformation( id="transformation:123", name="summary", title="Summary", description="Summarize the source", prompt="Summarize this", apply_default=False, model_id=model_id, created=datetime(2026, 1, 1, 12, 0, 0), updated=datetime(2026, 1, 1, 12, 0, 0), ) def _create_payload(model_id: str | None = None) -> dict: payload = { "name": "summary", "title": "Summary", "description": "Summarize the source", "prompt": "Summarize this", "apply_default": False, } if model_id is not None: payload["model_id"] = model_id return payload def test_create_transformation_with_model_id_persists_and_reads_back(): client = _client() saved_transformations: list[Transformation] = [] async def capture_save(transformation: Transformation): saved_transformations.append(transformation) transformation.id = "transformation:created" transformation.created = datetime(2026, 1, 1, 12, 0, 0) transformation.updated = datetime(2026, 1, 1, 12, 0, 0) with ( patch.object( Transformation, "save", autospec=True, side_effect=capture_save ), patch( "api.routers.transformations.Model.get", new_callable=AsyncMock, return_value=object(), ) as mock_model_get, ): response = client.post( "/api/transformations", json=_create_payload("model:local") ) assert response.status_code == 200 assert saved_transformations[0].model_id == "model:local" assert response.json()["model_id"] == "model:local" mock_model_get.assert_awaited_once_with("model:local") with patch( "api.routers.transformations.Transformation.get", new_callable=AsyncMock, return_value=saved_transformations[0], ): response = client.get("/api/transformations/transformation:created") assert response.status_code == 200 assert response.json()["model_id"] == "model:local" def test_execute_without_request_model_uses_stored_model(): client = _client() transformation = _transformation(model_id="model:stored") with ( patch( "api.routers.transformations.Transformation.get", new_callable=AsyncMock, return_value=transformation, ), patch( "api.routers.transformations.Model.get", new_callable=AsyncMock, return_value=object(), ) as mock_model_get, patch( "api.routers.transformations.transformation_graph.ainvoke", new_callable=AsyncMock, return_value={"output": "Stored model output"}, ) as mock_ainvoke, ): response = client.post( "/api/transformations/execute", json={ "transformation_id": transformation.id, "input_text": "Input text", }, ) assert response.status_code == 200 assert response.json()["model_id"] == "model:stored" mock_model_get.assert_awaited_once_with("model:stored") assert ( mock_ainvoke.call_args.kwargs["config"]["configurable"]["model_id"] == "model:stored" ) def test_execute_without_any_model_uses_default_transformation_model_fallback(): client = _client() transformation = _transformation(model_id=None) with ( patch( "api.routers.transformations.Transformation.get", new_callable=AsyncMock, return_value=transformation, ), patch( "api.routers.transformations.Model.get", new_callable=AsyncMock, ) as mock_model_get, patch( "api.routers.transformations.transformation_graph.ainvoke", new_callable=AsyncMock, return_value={"output": "Default model output"}, ) as mock_ainvoke, ): response = client.post( "/api/transformations/execute", json={ "transformation_id": transformation.id, "input_text": "Input text", }, ) assert response.status_code == 200 assert response.json()["model_id"] is None mock_model_get.assert_not_awaited() assert mock_ainvoke.call_args.kwargs["config"]["configurable"]["model_id"] is None def test_update_transformation_model_id_is_used_by_subsequent_execution(): client = _client() transformation = _transformation(model_id="model:old") async def save_update(transformation: Transformation): transformation.updated = datetime(2026, 1, 1, 12, 5, 0) with ( patch( "api.routers.transformations.Transformation.get", new_callable=AsyncMock, return_value=transformation, ), patch.object( Transformation, "save", autospec=True, side_effect=save_update ), patch( "api.routers.transformations.Model.get", new_callable=AsyncMock, return_value=object(), ) as mock_model_get, patch( "api.routers.transformations.transformation_graph.ainvoke", new_callable=AsyncMock, return_value={"output": "Updated model output"}, ) as mock_ainvoke, ): update_response = client.put( f"/api/transformations/{transformation.id}", json={"model_id": "model:new"}, ) execute_response = client.post( "/api/transformations/execute", json={ "transformation_id": transformation.id, "input_text": "Input text", }, ) assert update_response.status_code == 200 assert update_response.json()["model_id"] == "model:new" assert execute_response.status_code == 200 assert execute_response.json()["model_id"] == "model:new" # model:new is validated once at update time and looked up again at execute. assert mock_model_get.await_count == 2 mock_model_get.assert_awaited_with("model:new") assert ( mock_ainvoke.call_args.kwargs["config"]["configurable"]["model_id"] == "model:new" ) def test_execute_request_model_overrides_stored_model(): client = _client() transformation = _transformation(model_id="model:stored") with ( patch( "api.routers.transformations.Transformation.get", new_callable=AsyncMock, return_value=transformation, ), patch( "api.routers.transformations.Model.get", new_callable=AsyncMock, return_value=object(), ) as mock_model_get, patch( "api.routers.transformations.transformation_graph.ainvoke", new_callable=AsyncMock, return_value={"output": "Override output"}, ) as mock_ainvoke, ): response = client.post( "/api/transformations/execute", json={ "transformation_id": transformation.id, "input_text": "Input text", "model_id": "model:override", }, ) assert response.status_code == 200 assert response.json()["model_id"] == "model:override" mock_model_get.assert_awaited_once_with("model:override") assert ( mock_ainvoke.call_args.kwargs["config"]["configurable"]["model_id"] == "model:override" )