"""Tests for application/agents/workflow_agent.py graph loading and saving. Tests _parse_embedded_workflow, _load_from_database, and _finalize_workflow_run against the ephemeral ``pg_conn`` fixture. Agent construction is bypassed via ``__new__`` so we avoid the BaseAgent's LLM/tool wiring. """ from contextlib import contextmanager from datetime import datetime, timezone from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest def _make_agent(*, workflow_id=None, workflow=None, workflow_owner=None, decoded_token=None): """Construct a WorkflowAgent bypassing BaseAgent.__init__.""" from application.agents.workflow_agent import WorkflowAgent agent = WorkflowAgent.__new__(WorkflowAgent) agent.workflow_id = workflow_id agent.workflow_owner = workflow_owner agent._workflow_data = workflow agent._engine = None agent._run_persisted = False agent._bridge_error = None agent.decoded_token = decoded_token or {} return agent @contextmanager def _patch_db(conn): @contextmanager def _yield(): yield conn with patch( "application.agents.workflow_agent.db_readonly", _yield ), patch( "application.agents.workflow_agent.db_session", _yield ): yield class TestParseEmbeddedWorkflow: def test_returns_none_for_malformed_workflow(self): agent = _make_agent(workflow={"nodes": None}) # invalid shape assert agent._parse_embedded_workflow() is None def test_parses_valid_embedded_workflow(self): agent = _make_agent( workflow={ "name": "Embedded", "description": "d", "nodes": [ { "id": "n1", "type": "start", "title": "Start", "position": {"x": 0, "y": 0}, "data": {}, }, { "id": "n2", "type": "end", "title": "End", "position": {"x": 100, "y": 0}, "data": {}, }, ], "edges": [ { "id": "e1", "source": "n1", "target": "n2", } ], } ) graph = agent._parse_embedded_workflow() assert graph is not None assert graph.workflow.name == "Embedded" assert len(graph.nodes) == 2 assert len(graph.edges) == 1 def test_parses_edges_with_source_id_target_id(self): """Edges may arrive with ``source_id``/``target_id`` instead of ``source``/``target``.""" agent = _make_agent( workflow={ "name": "v2", "nodes": [ {"id": "a", "type": "start", "position": {"x": 0, "y": 0}, "data": {}}, {"id": "b", "type": "end", "position": {"x": 0, "y": 0}, "data": {}}, ], "edges": [ { "id": "e1", "source_id": "a", "target_id": "b", "source_handle": "h", "target_handle": "h2", } ], } ) graph = agent._parse_embedded_workflow() assert graph.edges[0].source_id == "a" class TestLoadWorkflowGraph: def test_returns_none_when_no_data_or_id(self): agent = _make_agent() assert agent._load_workflow_graph() is None def test_uses_embedded_when_provided(self): agent = _make_agent(workflow={"nodes": [], "edges": []}) # Empty nodes returns WorkflowGraph with no nodes, still valid object got = agent._load_workflow_graph() assert got is not None def test_uses_database_when_id_set(self, pg_conn): from application.storage.db.repositories.workflows import ( WorkflowsRepository, ) user = "u-loadwf" workflow = WorkflowsRepository(pg_conn).create(user, "wf") agent = _make_agent( workflow_id=str(workflow["id"]), workflow_owner=user, ) with _patch_db(pg_conn): got = agent._load_workflow_graph() assert got is not None class TestLoadFromDatabase: def test_returns_none_no_workflow_id(self): agent = _make_agent() assert agent._load_from_database() is None def test_returns_none_no_owner(self): agent = _make_agent(workflow_id="some-id") assert agent._load_from_database() is None def test_owner_from_decoded_token(self, pg_conn): from application.storage.db.repositories.workflows import ( WorkflowsRepository, ) user = "u-token-owner" wf = WorkflowsRepository(pg_conn).create(user, "wf") agent = _make_agent( workflow_id=str(wf["id"]), decoded_token={"sub": user}, ) with _patch_db(pg_conn): got = agent._load_from_database() assert got is not None def test_returns_none_when_workflow_missing(self, pg_conn): agent = _make_agent( workflow_id="00000000-0000-0000-0000-000000000000", workflow_owner="u", ) with _patch_db(pg_conn): got = agent._load_from_database() assert got is None def test_invalid_version_falls_back_to_1(self, pg_conn): """When current_graph_version is invalid it falls back to 1.""" from application.storage.db.repositories.workflows import ( WorkflowsRepository, ) user = "u-bad-version" wf = WorkflowsRepository(pg_conn).create(user, "wf") agent = _make_agent( workflow_id=str(wf["id"]), workflow_owner=user, ) with _patch_db(pg_conn): got = agent._load_from_database() assert got is not None def test_handles_exception(self): @contextmanager def _broken(): raise RuntimeError("boom") yield agent = _make_agent(workflow_id="x", workflow_owner="u") with patch( "application.agents.workflow_agent.db_readonly", _broken ): got = agent._load_from_database() assert got is None class TestSaveWorkflowRun: def test_returns_when_no_engine(self): agent = _make_agent(workflow_id="x", workflow_owner="u") # _engine is None agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "query") # should not raise def test_returns_when_no_workflow_id(self): agent = _make_agent(workflow_owner="u") agent._engine = MagicMock() agent._engine.execution_log = [] agent._engine.state = {} agent._engine.get_execution_summary.return_value = [] agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "query") def test_returns_when_workflow_missing_in_db(self, pg_conn): agent = _make_agent( workflow_id="00000000-0000-0000-0000-000000000000", workflow_owner="u", ) agent._engine = MagicMock() agent._engine.execution_log = [] agent._engine.state = {} agent._engine.get_execution_summary.return_value = [] with _patch_db(pg_conn): # Should just return None since workflow not found in DB agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "q") def test_creates_run_row(self, pg_conn): from application.storage.db.repositories.workflows import ( WorkflowsRepository, ) from application.storage.db.repositories.workflow_runs import ( WorkflowRunsRepository, ) import uuid as _uuid user = "u-saverun" wf = WorkflowsRepository(pg_conn).create(user, "wf") agent = _make_agent( workflow_id=str(wf["id"]), workflow_owner=user, ) agent._engine = MagicMock() agent._engine.execution_log = [] agent._engine.state = {"output": "hello"} agent._engine.get_execution_summary.return_value = [] # The run row is persisted under the engine's run id (the parent for any # run-scoped artifacts); a MagicMock id can't be adapted by the driver. run_id = str(_uuid.uuid4()) agent._engine.workflow_run_id = run_id with _patch_db(pg_conn): agent._finalize_workflow_run( agent.workflow_owner, agent.workflow_owner, None, "my query" ) runs = WorkflowRunsRepository(pg_conn).list_for_workflow(str(wf["id"])) assert len(runs) >= 1 assert str(runs[0]["id"]) == run_id def test_exception_is_swallowed(self): agent = _make_agent(workflow_id="x", workflow_owner="u") agent._engine = MagicMock() agent._engine.execution_log = [] agent._engine.state = {} agent._engine.get_execution_summary.return_value = [] @contextmanager def _broken(): raise RuntimeError("boom") yield with patch( "application.agents.workflow_agent.db_session", _broken ): # Should not raise agent._finalize_workflow_run(agent.workflow_owner, agent.workflow_owner, None, "q") class TestDetermineRunStatus: def test_completed_when_no_engine(self): from application.agents.workflows.schemas import ExecutionStatus agent = _make_agent() assert agent._determine_run_status() == ExecutionStatus.COMPLETED def test_completed_when_log_empty(self): from application.agents.workflows.schemas import ExecutionStatus agent = _make_agent() agent._engine = MagicMock() agent._engine.execution_log = [] assert agent._determine_run_status() == ExecutionStatus.COMPLETED def test_failed_if_any_log_failed(self): from application.agents.workflows.schemas import ExecutionStatus agent = _make_agent() agent._engine = MagicMock() agent._engine.execution_log = [ {"status": ExecutionStatus.COMPLETED.value}, {"status": ExecutionStatus.FAILED.value}, ] assert agent._determine_run_status() == ExecutionStatus.FAILED class TestSerializeState: def test_primitive_passes_through(self): agent = _make_agent() assert agent._serialize_state_value(42) == 42 assert agent._serialize_state_value("x") == "x" assert agent._serialize_state_value(None) is None assert agent._serialize_state_value(True) is True def test_datetime_becomes_iso(self): agent = _make_agent() now = datetime(2024, 1, 2, 3, 4, 5, tzinfo=timezone.utc) assert agent._serialize_state_value(now) == now.isoformat() def test_dict_becomes_dict_with_string_keys(self): agent = _make_agent() got = agent._serialize_state_value({1: "a", "b": 2}) assert got == {"1": "a", "b": 2} def test_tuple_becomes_list(self): agent = _make_agent() assert agent._serialize_state_value((1, 2, 3)) == [1, 2, 3] def test_unknown_becomes_string(self): agent = _make_agent() class Foo: def __str__(self): return "custom-foo" assert agent._serialize_state_value(Foo()) == "custom-foo" def test_serialize_state_dict(self): agent = _make_agent() got = agent._serialize_state({"x": 1, "y": [1, 2]}) assert got == {"x": 1, "y": [1, 2]} class TestGen: def test_yields_error_when_graph_fails(self): agent = _make_agent() # no workflow data # _load_workflow_graph returns None so gen yields error results = list(agent._gen_inner("q", log_context=None)) assert results == [ {"type": "error", "error": "Failed to load workflow configuration."} ] class TestAgentNodeApprovalPause: """A node agent whose tool pauses for approval must fail the node visibly, not emit empty output.""" def test_tool_calls_pending_raises_clear_error(self, monkeypatch): from application.agents.workflows import workflow_engine as we from application.agents.workflows.schemas import ( NodeType, Workflow, WorkflowGraph, WorkflowNode, ) from application.agents.workflows.workflow_engine import WorkflowEngine # An ephemeral node agent whose LLM handler yields the pause signal and ends, # emitting no "answer". Previously the engine dropped it and the node completed # with empty output (or raised a confusing "Structured output was expected"). class _PendingAgent: attachments = None def gen(self, prompt): yield {"type": "tool_calls_pending", "data": {"pending_tool_calls": [{"id": "x"}]}} monkeypatch.setattr( we.WorkflowNodeAgentFactory, "create", lambda **kw: _PendingAgent() ) graph = WorkflowGraph(workflow=Workflow(name="Approval Pause"), nodes=[], edges=[]) agent = SimpleNamespace( endpoint="stream", llm_name="openai", model_id="gpt-4o-mini", api_key="k", chat_history=[], user="u", decoded_token={"sub": "u"}, ) engine = WorkflowEngine( graph, agent, workflow_run_id="22222222-2222-2222-2222-222222222222" ) node = WorkflowNode( id="agent_1", workflow_id="wf-1", type=NodeType.AGENT, title="Extractor", position={"x": 0, "y": 0}, config={"llm_name": "openai", "system_prompt": "s", "tools": ["t"]}, ) with pytest.raises(ValueError, match="requires approval"): list(engine._execute_agent_node(node))