test_scheduler.py 16.7 KB
"""U-032 tests: the scheduler's own bookkeeping (loop-closing wiring + real
thread-level concurrency bookkeeping), mocked at the org_level_loop/
org_loop_route/run_goal_to_completion boundary -- this file is not about
re-proving e2e_pipeline/org_level_loop's own mechanics (covered in their own
test files), it's about proving finish_task_lifecycle applies the right
loop-closing step and run_scheduler_once actually runs goals concurrently
(not accidentally serially) while still collecting every result. Real
multi-Goal coverage (RUN_REAL_AGENT_E2E=1) is in src/test_scheduler_e2e.py.

Run: python3 -m unittest src.test_scheduler -v
"""

import os
import tempfile
import threading
import time
import unittest
from unittest.mock import patch

from src import org_level_loop, org_loop_route, schema, scheduler, settings_config, work_graph_execute, work_node_run


class FinishTaskLifecycleTestBase(unittest.TestCase):
    def setUp(self):
        self.conn = schema.connect(":memory:")
        schema.create_schema(self.conn)
        self.conn.execute(
            "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
            "VALUES ('proj-1', 'P', '/tmp/proj-1', 1, 'active')"
        )
        self.conn.execute(
            "INSERT INTO task (task_id, project_id, goal_text, task_type, status, loop_round, assigned_org_node) "
            "VALUES ('task-1', 'proj-1', 'goal', 'goal', 'validating', 0, 'validator')"
        )
        self.conn.commit()

    def tearDown(self):
        self.conn.close()


class FinishTaskLifecyclePassTest(FinishTaskLifecycleTestBase):
    def test_pass_verdict_closes_task_done(self):
        run_result = {
            "task_id": "task-1", "doc_artifact_ids": {}, "code_artifact_ids": {},
            "verdict": {"verdict": "PASS", "raw_response": "VERDICT: PASS"},
        }
        result = scheduler.finish_task_lifecycle(self.conn, run_result, "/tmp/proj-1", "goal text")
        self.assertEqual(result["final_status"], "done")
        status = self.conn.execute("SELECT status FROM task WHERE task_id = 'task-1'").fetchone()[0]
        self.assertEqual(status, "done")


class FinishTaskLifecycleFailTest(FinishTaskLifecycleTestBase):
    def test_fail_verdict_runs_fix_loop(self):
        run_result = {
            "task_id": "task-1",
            "doc_artifact_ids": {"benchmark": "artifact-b"},
            "code_artifact_ids": {},
            "verdict": {"verdict": "FAIL", "raw_response": "VERDICT: FAIL\nwall collision broken"},
        }
        with patch.object(scheduler.e2e_pipeline, "_read_artifact_content", return_value="benchmark content"), \
             patch.object(org_level_loop, "run_fix_loop", return_value={"final_status": "done"}) as mock_fix:
            result = scheduler.finish_task_lifecycle(self.conn, run_result, "/tmp/proj-1", "goal text")

        self.assertEqual(result["final_status"], "done")
        mock_fix.assert_called_once()
        args, kwargs = mock_fix.call_args
        self.assertEqual(args[0], self.conn)
        self.assertEqual(args[1], "task-1")
        self.assertEqual(args[2], "/tmp/proj-1")


class RunSchedulerOnceTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
        self.tmp.close()
        self.db_path = self.tmp.name
        conn = schema.connect(self.db_path)
        schema.create_schema(conn)
        conn.close()
        self.addCleanup(lambda: os.path.exists(self.db_path) and os.remove(self.db_path))

    def test_runs_all_goals_and_collects_results_in_order(self):
        def fake_run(db_path, target_folder_path, goal_text):
            return {"goal_text": goal_text, "final_status": "done"}

        goals = [("/tmp/p", "goal-{}".format(i)) for i in range(4)]
        with patch.object(scheduler, "run_goal_to_completion", side_effect=fake_run):
            results = scheduler.run_scheduler_once(self.db_path, goals, max_workers=2)

        self.assertEqual([r["goal_text"] for r in results], [g[1] for g in goals])
        self.assertTrue(all("result" in r for r in results))

    def test_one_goal_raising_does_not_stop_the_others(self):
        def fake_run(db_path, target_folder_path, goal_text):
            if goal_text == "goal-bad":
                raise RuntimeError("boom")
            return {"goal_text": goal_text, "final_status": "done"}

        goals = [("/tmp/p", "goal-ok-1"), ("/tmp/p", "goal-bad"), ("/tmp/p", "goal-ok-2")]
        with patch.object(scheduler, "run_goal_to_completion", side_effect=fake_run):
            results = scheduler.run_scheduler_once(self.db_path, goals, max_workers=3)

        self.assertIn("result", results[0])
        self.assertIn("error", results[1])
        self.assertEqual(results[1]["error"], "boom")
        self.assertIn("result", results[2])

    def test_max_workers_bounds_real_concurrency(self):
        current = [0]
        peak = [0]
        lock = threading.Lock()

        def fake_run(db_path, target_folder_path, goal_text):
            with lock:
                current[0] += 1
                peak[0] = max(peak[0], current[0])
            time.sleep(0.05)
            with lock:
                current[0] -= 1
            return {"goal_text": goal_text, "final_status": "done"}

        goals = [("/tmp/p", "goal-{}".format(i)) for i in range(6)]
        with patch.object(scheduler, "run_goal_to_completion", side_effect=fake_run):
            results = scheduler.run_scheduler_once(self.db_path, goals, max_workers=2)

        self.assertEqual(len(results), 6)
        self.assertLessEqual(peak[0], 2)
        self.assertEqual(peak[0], 2, "expected real parallelism up to the max_workers bound, not accidental serial execution")

    def test_defaults_max_workers_to_org_concurrency_config(self):
        conn = schema.connect(self.db_path)
        settings_config.upsert_concurrency_config(conn, max_concurrent_goals=3, updated_by="test")
        conn.close()

        current = [0]
        peak = [0]
        lock = threading.Lock()

        def fake_run(db_path, target_folder_path, goal_text):
            with lock:
                current[0] += 1
                peak[0] = max(peak[0], current[0])
            time.sleep(0.05)
            with lock:
                current[0] -= 1
            return {"goal_text": goal_text, "final_status": "done"}

        goals = [("/tmp/p", "goal-{}".format(i)) for i in range(6)]
        with patch.object(scheduler, "run_goal_to_completion", side_effect=fake_run):
            scheduler.run_scheduler_once(self.db_path, goals, max_workers=None)

        self.assertEqual(peak[0], 3)


class ConcurrentSqliteWritesStressTest(unittest.TestCase):
    """Not gated behind RUN_REAL_AGENT_E2E -- doesn't call the real Claude
    API, but DOES exercise the actual novel risk in this module's design
    (see module docstring's "Threading model" note): multiple threads, each
    holding its OWN sqlite connection, writing concurrently to the SAME db
    file. e2e_pipeline.run_end_to_end is faked (no real API calls) but
    finish_task_lifecycle's own close_task_done still runs for real against
    a real connection, so this is a genuine concurrent-write stress test,
    not a fully mocked one."""

    def setUp(self):
        self.tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
        self.tmp.close()
        self.db_path = self.tmp.name
        conn = schema.connect(self.db_path)
        schema.create_schema(conn)
        conn.execute(
            "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
            "VALUES ('proj-1', 'P', '/tmp/proj-1', 1, 'active')"
        )
        conn.commit()
        conn.close()
        self.addCleanup(lambda: os.path.exists(self.db_path) and os.remove(self.db_path))

    def test_many_concurrent_goals_write_to_the_same_db_file_without_locking_errors(self):
        def fake_run_task_lifecycle(conn, task_id, target_folder_path, goal_text):
            # task_id already exists -- run_goal_to_completion creates it for
            # real via task_create.create_task before calling this; stand in
            # for org_node_lifecycle's real queued->...->validating march so
            # finish_task_lifecycle's close_task_done (requires 'validating') works.
            conn.execute("UPDATE task SET status = 'validating' WHERE task_id = ?", (task_id,))
            conn.commit()
            time.sleep(0.02)  # widen the contention window
            return {
                "task_id": task_id, "doc_artifact_ids": {}, "code_artifact_ids": {},
                "verdict": {"verdict": "PASS", "raw_response": "VERDICT: PASS"},
            }

        goals = [("/tmp/proj-1", "goal-{}".format(i)) for i in range(10)]
        with patch.object(scheduler.e2e_pipeline, "run_task_lifecycle", side_effect=fake_run_task_lifecycle):
            results = scheduler.run_scheduler_once(self.db_path, goals, max_workers=5)

        self.assertTrue(all("result" in r for r in results), results)

        conn = schema.connect(self.db_path)
        done_count = conn.execute("SELECT COUNT(*) FROM task WHERE status = 'done'").fetchone()[0]
        conn.close()
        self.assertEqual(done_count, 10)


class RoutableExceptionTestBase(unittest.TestCase):
    """REQ-06/07/14 failure: work-node-level failures (CorruptPersistedState
    /ClassifierBlocked/UsageLimitWaitTimeout) must actually reach
    task.status='blocked_pending_human' with the matching blocked_reason --
    found missing during the PH4-closure audit (2026-08-02): each exception
    was correctly raised and correctly marked its own work_node_execution
    'blocked', but nothing ever caught it to apply the task-level
    consequence 06's own acceptance text names."""

    def setUp(self):
        self.tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
        self.tmp.close()
        self.db_path = self.tmp.name
        conn = schema.connect(self.db_path)
        schema.create_schema(conn)
        conn.execute(
            "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
            "VALUES ('proj-1', 'P', '/tmp/proj-1', 1, 'active')"
        )
        conn.commit()
        conn.close()
        self.addCleanup(lambda: os.path.exists(self.db_path) and os.remove(self.db_path))

    def _run_with_raising_lifecycle(self, exc):
        def fake_run_task_lifecycle(conn, task_id, target_folder_path, goal_text):
            raise exc

        with patch.object(scheduler.e2e_pipeline, "run_task_lifecycle", side_effect=fake_run_task_lifecycle):
            result = scheduler.run_goal_to_completion(self.db_path, "/tmp/proj-1", "goal text")
        return result

    def _assert_task_blocked(self, task_id, blocked_reason):
        conn = schema.connect(self.db_path)
        row = conn.execute("SELECT status, blocked_reason FROM task WHERE task_id = ?", (task_id,)).fetchone()
        event = conn.execute(
            "SELECT event_type FROM execution_event WHERE task_id = ? ORDER BY rowid DESC LIMIT 1", (task_id,)
        ).fetchone()
        conn.close()
        self.assertEqual(row, ("blocked_pending_human", blocked_reason))
        self.assertEqual(event[0], "task_status_changed")


class CorruptPersistedStateRoutingTest(RoutableExceptionTestBase):
    def test_routes_to_blocked_pending_human(self):
        result = self._run_with_raising_lifecycle(work_node_run.CorruptPersistedStateError("boom"))
        self.assertEqual(result["final_status"], "blocked_pending_human")
        self.assertEqual(result["blocked_reason"], "corrupt_persisted_state")
        self._assert_task_blocked(result["task_id"], "corrupt_persisted_state")


class ClassifierBlockedRoutingTest(RoutableExceptionTestBase):
    def test_routes_to_blocked_pending_human(self):
        result = self._run_with_raising_lifecycle(work_graph_execute.ClassifierBlockedError("boom"))
        self.assertEqual(result["final_status"], "blocked_pending_human")
        self.assertEqual(result["blocked_reason"], "classifier_blocked")
        self._assert_task_blocked(result["task_id"], "classifier_blocked")


class UsageLimitWaitTimeoutRoutingTest(RoutableExceptionTestBase):
    def test_routes_to_blocked_pending_human(self):
        result = self._run_with_raising_lifecycle(work_node_run.UsageLimitWaitTimeoutError("boom"))
        self.assertEqual(result["final_status"], "blocked_pending_human")
        self.assertEqual(result["blocked_reason"], "usage_limit_wait_timeout")
        self._assert_task_blocked(result["task_id"], "usage_limit_wait_timeout")


class ConsistencyCheckFailedRoutingTest(RoutableExceptionTestBase):
    def test_routes_to_blocked_pending_human(self):
        result = self._run_with_raising_lifecycle(
            scheduler.work_graph_orchestrate.ConsistencyCheckFailedError("boom")
        )
        self.assertEqual(result["final_status"], "blocked_pending_human")
        self.assertEqual(result["blocked_reason"], "prd_consistency_check_failed")
        self._assert_task_blocked(result["task_id"], "prd_consistency_check_failed")


class UnroutedExceptionStillPropagatesTest(RoutableExceptionTestBase):
    def test_other_exceptions_are_not_swallowed(self):
        with self.assertRaises(RuntimeError):
            self._run_with_raising_lifecycle(RuntimeError("something else entirely"))


def _fix_loop_trigger(org_node, prompt, target_folder_path, timeout=300):
    """org_node_trigger.trigger_org_node fake for org_level_loop.run_fix_loop's
    3 perspective calls (all agree 'coder') + fixer + verify -- verify always
    fails, so route_org_loop keeps incrementing loop_round on retries."""
    if "CATEGORY:" in prompt:
        return "CATEGORY: coder\nDETAIL: because reasons"
    if org_node == "coder":
        return "Fixed it."
    if org_node == "validator":
        return "VERDICT: FAIL\nstill broken"
    raise AssertionError("unexpected call: " + org_node)


class SchedulerContinuesPastBlockedTaskTest(unittest.TestCase):
    """REQ-05 failure / REQ-09 boundary, against the REAL scheduler (U-032) --
    test_org_loop_route.py's own test for this literally says "no scheduler
    exists yet to test against directly" (stale as of U-032). One Goal's
    task is pre-seeded at loop_round=3 so its fix-loop attempt hits
    org_loop_route.LoopLimitReachedError for real (route_org_loop's own real
    blocked_pending_human write, not a mock) while the others complete via
    close_task_done -- proving the scheduler doesn't stop serving the rest
    of the batch."""

    def setUp(self):
        self.tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
        self.tmp.close()
        self.db_path = self.tmp.name
        conn = schema.connect(self.db_path)
        schema.create_schema(conn)
        conn.execute(
            "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
            "VALUES ('proj-1', 'P', '/tmp/proj-1', 1, 'active')"
        )
        conn.commit()
        conn.close()
        self.addCleanup(lambda: os.path.exists(self.db_path) and os.remove(self.db_path))

    def test_one_task_hitting_loop_limit_does_not_block_the_others(self):
        def fake_run_task_lifecycle(conn, task_id, target_folder_path, goal_text):
            conn.execute("UPDATE task SET status = 'validating' WHERE task_id = ?", (task_id,))
            conn.commit()
            if goal_text == "goal-doomed":
                conn.execute("UPDATE task SET loop_round = 3 WHERE task_id = ?", (task_id,))
                conn.commit()
                verdict = {"verdict": "FAIL", "raw_response": "VERDICT: FAIL\nstill broken"}
            else:
                verdict = {"verdict": "PASS", "raw_response": "VERDICT: PASS"}
            return {
                "task_id": task_id, "doc_artifact_ids": {"benchmark": "artifact-b"},
                "code_artifact_ids": {}, "verdict": verdict,
            }

        goals = [("/tmp/proj-1", "goal-ok-1"), ("/tmp/proj-1", "goal-doomed"), ("/tmp/proj-1", "goal-ok-2")]
        with patch.object(scheduler.e2e_pipeline, "run_task_lifecycle", side_effect=fake_run_task_lifecycle), \
             patch.object(scheduler.e2e_pipeline, "_read_artifact_content", return_value="benchmark content"), \
             patch("src.org_node_trigger.trigger_org_node", side_effect=_fix_loop_trigger):
            results = scheduler.run_scheduler_once(self.db_path, goals, max_workers=3)

        by_goal = {r["goal_text"]: r for r in results}
        self.assertEqual(by_goal["goal-ok-1"]["result"]["final_status"], "done")
        self.assertEqual(by_goal["goal-ok-2"]["result"]["final_status"], "done")
        self.assertIn("error", by_goal["goal-doomed"])

        conn = schema.connect(self.db_path)
        doomed_status = conn.execute(
            "SELECT status, blocked_reason FROM task WHERE goal_text = 'goal-doomed'"
        ).fetchone()
        conn.close()
        self.assertEqual(doomed_status, ("blocked_pending_human", "loop_limit_reached"))


if __name__ == "__main__":
    unittest.main()