from __future__ import annotations from unittest import IsolatedAsyncioTestCase from unittest.mock import AsyncMock, patch from uuid import UUID from .contracts.engine_gateway import EngineMessage, GenerateRequest from .engine_client import EngineClient from .services import continuous_improvement_agentic as agentic from .services import continuous_improvement_producer as producer JOB_ID = UUID("81000000-0000-0000-0000-000000000001") def _claimed_job() -> producer.ClaimedAgenticJob: spec = producer.load_repo_approved_job() return producer.ClaimedAgenticJob( job_id=JOB_ID, spec=spec, source_fingerprint=producer._fingerprint(spec), attempt_count=1, ) def _pipeline_result(*, replay: bool = False) -> agentic.AgenticPipelineResult: ids = producer._job_ids(JOB_ID) return agentic.AgenticPipelineResult( submission_id=ids["submission"], pipeline_id=ids["pipeline"], qualification_id=ids["qualification"], candidate_catalog_entry_id="oas-g8-catalog-scheduled-test", state="pending_human_approval", human_approval_required=True, catalog_promoted=False, idempotent_replay=replay, clinical_claim_allowed=False, draft_id="oas-g8-draft-scheduled-test", benchmark_id="oas-g8-benchmark-scheduled-test", red_team_review_count=2, benchmark_variant_count=3, agent_calls_executed=0 if replay else 7, trigger_kind="scheduled_repo_source", ) class AcquireContext: def __init__(self, conn: AsyncMock) -> None: self.conn = conn async def __aenter__(self) -> AsyncMock: return self.conn async def __aexit__(self, *_: object) -> None: return None class ContinuousImprovementProducerTest(IsolatedAsyncioTestCase): def test_repo_source_is_approved_hashed_and_synthetic_only(self) -> None: spec = producer.load_repo_approved_job() self.assertEqual(spec.data_classification, producer.DATA_CLASSIFICATION) self.assertEqual(spec.trigger_kind, "scheduled_repo_source") self.assertEqual(spec.variant_count, 3) self.assertTrue(all(item.artifact.usage_status == "approved" for item in spec.source_packs)) self.assertEqual(len(producer._fingerprint(spec)), 64) async def test_restricted_source_is_rejected_before_enqueue_sql(self) -> None: payload = producer.load_repo_approved_job().model_dump(mode="json") payload["source_packs"][0]["artifact"]["usage_status"] = "restricted" with self.assertRaises(ValueError): producer.ScheduledAgenticJobSpec.model_validate(payload) async def test_success_only_marks_durable_pending_human_candidate_complete(self) -> None: conn = AsyncMock() run = AsyncMock(return_value=_pipeline_result()) with ( patch.object(producer.db, "acquire", return_value=AcquireContext(conn)), patch.object(agentic, "run_agentic_content_pipeline", run), ): outcome = await producer.execute_claimed_agentic_job(_claimed_job()) self.assertEqual(outcome.status, "completed") self.assertEqual(outcome.agent_calls_executed, 7) update_sql = str(conn.execute.await_args.args[0]) self.assertIn("status = 'completed'", update_sql) self.assertNotIn("ci_catalog_entry", update_sql) kwargs = run.await_args.kwargs self.assertEqual(kwargs["trigger_kind"], "scheduled_repo_source") self.assertEqual(kwargs["variant_count"], 3) async def test_engine_failure_rolls_back_candidate_and_leaves_retry_state(self) -> None: conn = AsyncMock() mark_retry = AsyncMock() run = AsyncMock(side_effect=agentic.AgenticPipelineExecutionError("engine down")) with ( patch.object(producer.db, "acquire", return_value=AcquireContext(conn)), patch.object(agentic, "run_agentic_content_pipeline", run), patch.object(producer, "_mark_retry", mark_retry), ): outcome = await producer.execute_claimed_agentic_job(_claimed_job()) self.assertEqual(outcome.status, "retry_wait") mark_retry.assert_awaited_once() self.assertEqual(conn.execute.await_count, 0) async def test_safety_rejection_never_marks_result_complete(self) -> None: conn = AsyncMock() mark_rejected = AsyncMock() run = AsyncMock(side_effect=agentic.AgenticPipelineRejectedError("PII")) with ( patch.object(producer.db, "acquire", return_value=AcquireContext(conn)), patch.object(agentic, "run_agentic_content_pipeline", run), patch.object(producer, "_mark_rejected", mark_rejected), ): outcome = await producer.execute_claimed_agentic_job(_claimed_job()) self.assertEqual(outcome.status, "rejected") mark_rejected.assert_awaited_once() self.assertEqual(conn.execute.await_count, 0) async def test_idempotent_recovery_completes_without_model_calls(self) -> None: conn = AsyncMock() run = AsyncMock(return_value=_pipeline_result(replay=True)) with ( patch.object(producer.db, "acquire", return_value=AcquireContext(conn)), patch.object(agentic, "run_agentic_content_pipeline", run), ): outcome = await producer.execute_claimed_agentic_job(_claimed_job()) self.assertEqual(outcome.status, "completed") self.assertTrue(outcome.idempotent_replay) self.assertEqual(outcome.agent_calls_executed, 0) async def test_cycle_isolates_retry_and_continues_other_jobs(self) -> None: jobs = [_claimed_job(), _claimed_job(), _claimed_job()] outcomes = [ producer.AgenticJobOutcome( job_id=JOB_ID, status="retry_wait", agent_calls_executed=0, error_code="engine_execution_failed", ), producer.AgenticJobOutcome( job_id=JOB_ID, status="completed", agent_calls_executed=7, ), producer.AgenticJobOutcome( job_id=JOB_ID, status="rejected", agent_calls_executed=0, error_code="safety_gate_rejected", ), ] with ( patch.object(producer, "ensure_repo_approved_job", AsyncMock(return_value=JOB_ID)), patch.object(producer, "claim_next_agentic_job", AsyncMock(side_effect=jobs)), patch.object(producer, "execute_claimed_agentic_job", AsyncMock(side_effect=outcomes)), patch.object(producer.settings, "continuous_improvement_producer_batch_size", 3), ): result = await producer.produce_queued_agentic_jobs_once() self.assertEqual(result["claimed"], 3) self.assertEqual(result["completed"], 1) self.assertEqual(result["retry_wait"], 1) self.assertEqual(result["rejected"], 1) async def test_scheduler_is_disabled_by_default_setting(self) -> None: with patch.object( producer.settings, "continuous_improvement_producer_enabled", False, ): self.assertIsNone(producer.schedule_continuous_improvement_producer()) async def test_producer_uses_its_explicit_engine_timeout(self) -> None: engine = EngineClient(base_url="http://engine.test") request = GenerateRequest( ai_role="evaluator", messages=[EngineMessage(role="user", content="synthetic")], ) with patch.object(engine, "generate", AsyncMock()) as generate: await producer._ProducerEngine(engine).generate(request) self.assertEqual( generate.await_args.kwargs["timeout"], producer.settings.continuous_improvement_producer_engine_timeout_seconds, ) if __name__ == "__main__": import unittest unittest.main()