Текстовый файл из учебного комплекта, не свидетельство установленного runtime. Вернуться к учебнику · Исходный файл
1 import unittest 2 from unittest.mock import AsyncMock, patch 3 from uuid import uuid4 4 5 from src_agent.agent import AgentRunner 6 from src_agent.repositories.runtime_state import RuntimeStateConflict 7 from src_shared.contracts import ( 8 AgentInboundRequest, 9 WorkflowRunResponse, 10 WorkflowRuntimeState, 11 ) 12 13 14 class _Transaction: 15 def __init__(self, connection): 16 self.connection = connection 17 self.exception = None 18 19 async def __aenter__(self): 20 return self.connection 21 22 async def __aexit__(self, exc_type, exc, traceback): 23 self.exception = exc 24 return False 25 26 27 class _Database: 28 def __init__(self): 29 self.connection = object() 30 self.transactions = [] 31 32 def transaction(self): 33 transaction = _Transaction(self.connection) 34 self.transactions.append(transaction) 35 return transaction 36 37 38 class AgentPersistenceTest(unittest.IsolatedAsyncioTestCase): 39 @staticmethod 40 def _request() -> AgentInboundRequest: 41 return AgentInboundRequest( 42 trace_id=uuid4(), 43 request_id=uuid4(), 44 session_id=None, 45 ts_ms=1, 46 text="hello", 47 ) 48 49 @staticmethod 50 def _response(request: AgentInboundRequest) -> WorkflowRunResponse: 51 return WorkflowRunResponse( 52 trace_id=request.trace_id, 53 request_id=request.request_id, 54 session_id=request.session_id, 55 ts_ms=request.ts_ms, 56 status="DONE", 57 result="world", 58 next_runtime=WorkflowRuntimeState(version=2), 59 ) 60 61 def _runner(self, database: _Database, response: WorkflowRunResponse) -> AgentRunner: 62 runner = AgentRunner( 63 nc=object(), 64 db=database, 65 workflow_subject="workflow.run", 66 user_id="test-user", 67 workflow_timeout_s=1, 68 ) 69 runner._call_workflow = AsyncMock(return_value=response) 70 return runner 71 72 async def test_runtime_state_and_history_share_one_transaction(self) -> None: 73 request = self._request() 74 response = self._response(request) 75 database = _Database() 76 runner = self._runner(database, response) 77 78 save_state = AsyncMock(return_value=2) 79 persist_turn = AsyncMock() 80 with ( 81 patch("src_agent.agent.save_runtime_state", save_state), 82 patch("src_agent.agent.persist_chat_turn", persist_turn), 83 ): 84 result = await runner._run_with_optimistic_lock( 85 req=request, 86 session_id=uuid4(), 87 runtime=WorkflowRuntimeState(version=1), 88 ) 89 90 self.assertEqual("DONE", result.status) 91 self.assertEqual(1, len(database.transactions)) 92 self.assertIs(database.connection, save_state.await_args.kwargs["connection"]) 93 self.assertIs(database.connection, persist_turn.await_args.kwargs["connection"]) 94 95 async def test_persistence_failure_does_not_repeat_workflow_side_effects(self) -> None: 96 request = self._request() 97 response = self._response(request) 98 database = _Database() 99 runner = self._runner(database, response) 100 101 with ( 102 patch("src_agent.agent.save_runtime_state", AsyncMock(return_value=2)), 103 patch( 104 "src_agent.agent.persist_chat_turn", 105 AsyncMock(side_effect=RuntimeError("database unavailable")), 106 ), 107 self.assertRaisesRegex(RuntimeError, "database unavailable"), 108 ): 109 await runner._run_with_optimistic_lock( 110 req=request, 111 session_id=uuid4(), 112 runtime=WorkflowRuntimeState(version=1), 113 ) 114 115 self.assertEqual(1, runner._call_workflow.await_count) 116 self.assertIsInstance(database.transactions[0].exception, RuntimeError) 117 118 async def test_only_version_conflicts_repeat_the_workflow(self) -> None: 119 request = self._request() 120 response = self._response(request) 121 database = _Database() 122 runner = self._runner(database, response) 123 refreshed = WorkflowRuntimeState(version=2) 124 125 save_state = AsyncMock( 126 side_effect=[RuntimeStateConflict("conflict"), 3], 127 ) 128 with ( 129 patch("src_agent.agent.save_runtime_state", save_state), 130 patch("src_agent.agent.persist_chat_turn", AsyncMock()), 131 patch( 132 "src_agent.agent.load_runtime_state", 133 AsyncMock(return_value=refreshed), 134 ) as load_state, 135 ): 136 result = await runner._run_with_optimistic_lock( 137 req=request, 138 session_id=uuid4(), 139 runtime=WorkflowRuntimeState(version=1), 140 ) 141 142 self.assertEqual("DONE", result.status) 143 self.assertEqual(2, runner._call_workflow.await_count) 144 self.assertEqual(1, load_state.await_count) 145 146 147 if __name__ == "__main__": 148 unittest.main()