宜博 / lemongraph

"""U-001 tests: all tables build, and the constraints named in the Unit's
Done criterion (unique / foreign key / enum) are actually enforced.
U-034 (PH5, 2026-08-02) added `deployment`/`deploy_health_signal` (17 tables
total) and widened `bug_record.root_cause_category` to accept `core_engine`.

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

import sqlite3
import unittest

from src import schema


class SchemaCreationTest(unittest.TestCase):
    def setUp(self):
        self.conn = schema.connect(":memory:")
        schema.create_schema(self.conn)

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

    def test_all_17_tables_created(self):
        rows = self.conn.execute(
            "SELECT name FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%'"
        ).fetchall()
        actual = {row[0] for row in rows}
        self.assertEqual(actual, set(schema.TABLE_NAMES))
        self.assertEqual(len(schema.TABLE_NAMES), 17)

    def test_schema_is_idempotent_to_reopen_on_disk(self):
        # not required by the Done criterion, but a schema that only works once
        # against :memory: would be a trap for every later Unit that opens a
        # real db file -- cheap to check now.
        import os
        import tempfile

        fd, path = tempfile.mkstemp(suffix=".sqlite3")
        os.close(fd)
        os.remove(path)
        try:
            conn = schema.connect(path)
            schema.create_schema(conn)
            conn.close()
            conn2 = schema.connect(path)
            rows = conn2.execute(
                "SELECT name FROM sqlite_master WHERE type = 'table' AND name NOT LIKE 'sqlite_%'"
            ).fetchall()
            conn2.close()
            self.assertEqual({r[0] for r in rows}, set(schema.TABLE_NAMES))
        finally:
            if os.path.exists(path):
                os.remove(path)


class ConstraintTestBase(unittest.TestCase):
    def setUp(self):
        self.conn = schema.connect(":memory:")
        schema.create_schema(self.conn)

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

    def insert_project(self, project_id="proj-1", target_folder_path="/tmp/proj-1"):
        self.conn.execute(
            "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
            "VALUES (?, ?, ?, 1, 'active')",
            (project_id, "Project One", target_folder_path),
        )
        self.conn.commit()
        return project_id

    def insert_task(self, task_id="task-1", project_id="proj-1", status="queued"):
        self.conn.execute(
            "INSERT INTO task (task_id, project_id, goal_text, task_type, status, assigned_org_node) "
            "VALUES (?, ?, 'implement snake game', 'goal', ?, 'planner')",
            (task_id, project_id, status),
        )
        self.conn.commit()
        return task_id

    def insert_skill(self, skill_id="skill-1"):
        self.conn.execute(
            "INSERT INTO skill (skill_id, owner_org_node, file_path, current_git_commit_sha, source, status) "
            "VALUES (?, 'coder', 'skills/skill-1/', 'abc123', 'new', 'active')",
            (skill_id,),
        )
        self.conn.commit()
        return skill_id

    def insert_work_graph_run(self, run_id="run-1", task_id="task-1", org_node="coder", status="planning"):
        self.conn.execute(
            "INSERT INTO work_graph_run (run_id, task_id, org_node, graph_manifest, status) "
            "VALUES (?, ?, ?, '{}', ?)",
            (run_id, task_id, org_node, status),
        )
        self.conn.commit()
        return run_id


class ProjectUniqueConstraintTest(ConstraintTestBase):
    def test_duplicate_target_folder_path_rejected(self):
        self.insert_project("proj-1", "/tmp/shared-path")
        with self.assertRaises(sqlite3.IntegrityError):
            self.conn.execute(
                "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
                "VALUES ('proj-2', 'Project Two', '/tmp/shared-path', 1, 'active')"
            )

    def test_has_git_repo_boolean_check_rejects_out_of_range_value(self):
        with self.assertRaises(sqlite3.IntegrityError):
            self.conn.execute(
                "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
                "VALUES ('proj-1', 'P', '/tmp/p', 2, 'active')"
            )


class WorkNodeExecutionEnumConstraintTest(ConstraintTestBase):
    def setUp(self):
        super().setUp()
        self.insert_project()
        self.insert_task()
        self.insert_skill()
        self.insert_work_graph_run()

    def _insert_node_exec(self, status):
        self.conn.execute(
            "INSERT INTO work_node_execution "
            "(node_exec_id, run_id, node_id, subagent_type, skill_ref, "
            " input_schema_ref, output_schema_ref, input_payload, status) "
            "VALUES ('node-1', 'run-1', 'n1', 'coder', 'skill-1', 'in.json', 'out.json', '{}', ?)",
            (status,),
        )

    def test_valid_status_accepted(self):
        self._insert_node_exec("waiting_for_quota")
        self.conn.commit()
        row = self.conn.execute(
            "SELECT status FROM work_node_execution WHERE node_exec_id = 'node-1'"
        ).fetchone()
        self.assertEqual(row[0], "waiting_for_quota")

    def test_invalid_status_rejected(self):
        with self.assertRaises(sqlite3.IntegrityError):
            self._insert_node_exec("not_a_real_status")


class ForeignKeyEnforcementTest(ConstraintTestBase):
    def test_task_rejects_unknown_project_id(self):
        with self.assertRaises(sqlite3.IntegrityError):
            self.insert_task(task_id="task-1", project_id="does-not-exist")

    def test_work_node_execution_rejects_unknown_skill_ref(self):
        self.insert_project()
        self.insert_task()
        self.insert_work_graph_run()
        with self.assertRaises(sqlite3.IntegrityError):
            self.conn.execute(
                "INSERT INTO work_node_execution "
                "(node_exec_id, run_id, node_id, subagent_type, skill_ref, "
                " input_schema_ref, output_schema_ref, input_payload, status) "
                "VALUES ('node-1', 'run-1', 'n1', 'coder', 'no-such-skill', 'in.json', 'out.json', '{}', 'pending')"
            )


class WorkGraphRunActiveUniqueIndexTest(ConstraintTestBase):
    """01-M 1.3: only one active (planning/running/merging/verifying) run per
    (task_id, org_node) at a time -- terminal-status runs don't count."""

    def setUp(self):
        super().setUp()
        self.insert_project()
        self.insert_task()

    def test_second_active_run_for_same_task_and_org_node_rejected(self):
        self.insert_work_graph_run("run-1", "task-1", "coder", status="running")
        with self.assertRaises(sqlite3.IntegrityError):
            self.insert_work_graph_run("run-2", "task-1", "coder", status="planning")

    def test_second_run_allowed_once_first_is_terminal(self):
        self.insert_work_graph_run("run-1", "task-1", "coder", status="done")
        self.insert_work_graph_run("run-2", "task-1", "coder", status="planning")
        count = self.conn.execute(
            "SELECT COUNT(*) FROM work_graph_run WHERE task_id = 'task-1' AND org_node = 'coder'"
        ).fetchone()[0]
        self.assertEqual(count, 2)


class ArtifactUniqueConstraintTest(ConstraintTestBase):
    def setUp(self):
        super().setUp()
        self.insert_project()
        self.insert_task()
        self.insert_work_graph_run()

    def _insert_artifact(self, artifact_id):
        self.conn.execute(
            "INSERT INTO artifact (artifact_id, artifact_type, task_id, produced_by_run_id, "
            "file_path, git_commit_sha) VALUES (?, 'sdd', 'task-1', 'run-1', 'sdd.md', 'sha-1')",
            (artifact_id,),
        )

    def test_duplicate_task_type_commit_rejected(self):
        self._insert_artifact("artifact-1")
        self.conn.commit()
        with self.assertRaises(sqlite3.IntegrityError):
            self._insert_artifact("artifact-2")

    def test_new_test_artifact_types_accepted(self):
        for artifact_type in (
            "test_dataset", "test_script",
            "validator_benchmark", "validator_test_dataset", "validator_test_script",
        ):
            self.conn.execute(
                "INSERT INTO artifact (artifact_id, artifact_type, task_id, produced_by_run_id, "
                "file_path, git_commit_sha) VALUES (?, ?, 'task-1', 'run-1', 'x.md', 'sha-1')",
                ("artifact-" + artifact_type, artifact_type),
            )
        self.conn.commit()


class MigrateExpandArtifactTypesTest(unittest.TestCase):
    """2026-08-06: real upgrade path for an already-populated database whose
    artifact table still has the pre-expansion CHECK constraint."""

    def setUp(self):
        self.conn = schema.connect(":memory:")
        # Simulate a pre-migration database: create the OLD schema by hand
        # (not schema.create_schema(), which already has the new constraint).
        old_sql = schema.SCHEMA_SQL.replace(
            "'brd', 'prd', 'sdd', 'benchmark', 'bug_record', 'code_diff', 'test_report',\n"
            "        'test_dataset', 'test_script',\n"
            "        'validator_benchmark', 'validator_test_dataset', 'validator_test_script'",
            "'brd', 'prd', 'sdd', 'benchmark', 'bug_record', 'code_diff', 'test_report'",
        )
        self.conn.executescript(old_sql)
        self.conn.commit()
        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, assigned_org_node) "
            "VALUES ('task-1', 'proj-1', 'g', 'goal', 'queued', 'planner')"
        )
        self.conn.execute(
            "INSERT INTO work_graph_run (run_id, task_id, org_node, graph_manifest, status) "
            "VALUES ('run-1', 'task-1', 'planner', '{}', 'planning')"
        )
        self.conn.execute(
            "INSERT INTO artifact (artifact_id, artifact_type, task_id, produced_by_run_id, "
            "file_path, git_commit_sha) VALUES ('art-1', 'brd', 'task-1', 'run-1', 'brd.md', 'sha-1')"
        )
        self.conn.commit()

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

    def test_old_check_constraint_rejects_new_type_before_migration(self):
        with self.assertRaises(sqlite3.IntegrityError):
            self.conn.execute(
                "INSERT INTO artifact (artifact_id, artifact_type, task_id, produced_by_run_id, "
                "file_path, git_commit_sha) VALUES ('art-2', 'test_script', 'task-1', 'run-1', 'x', 'sha-2')"
            )

    def test_migration_preserves_existing_rows_and_accepts_new_types(self):
        schema.migrate_expand_artifact_types(self.conn)

        row = self.conn.execute("SELECT artifact_type FROM artifact WHERE artifact_id = 'art-1'").fetchone()
        self.assertEqual(row[0], "brd")

        self.conn.execute(
            "INSERT INTO artifact (artifact_id, artifact_type, task_id, produced_by_run_id, "
            "file_path, git_commit_sha) VALUES ('art-2', 'test_script', 'task-1', 'run-1', 'x', 'sha-2')"
        )
        self.conn.commit()
        count = self.conn.execute("SELECT COUNT(*) FROM artifact").fetchone()[0]
        self.assertEqual(count, 2)

    def test_unique_index_still_enforced_after_migration(self):
        schema.migrate_expand_artifact_types(self.conn)
        self.conn.commit()
        with self.assertRaises(sqlite3.IntegrityError):
            self.conn.execute(
                "INSERT INTO artifact (artifact_id, artifact_type, task_id, produced_by_run_id, "
                "file_path, git_commit_sha) VALUES ('art-dup', 'brd', 'task-1', 'run-1', 'brd.md', 'sha-1')"
            )

    def test_idempotent_on_a_database_already_migrated(self):
        schema.migrate_expand_artifact_types(self.conn)
        schema.migrate_expand_artifact_types(self.conn)  # must not raise / must not duplicate
        count = self.conn.execute("SELECT COUNT(*) FROM artifact").fetchone()[0]
        self.assertEqual(count, 1)

    def test_noop_on_a_freshly_created_already_expanded_database(self):
        fresh = schema.connect(":memory:")
        try:
            schema.create_schema(fresh)
            schema.migrate_expand_artifact_types(fresh)  # must not raise
            fresh.execute(
                "INSERT INTO project (project_id, name, target_folder_path, has_git_repo, status) "
                "VALUES ('p', 'P', '/tmp/p', 1, 'active')"
            )
            fresh.execute(
                "INSERT INTO task (task_id, project_id, goal_text, task_type, status, assigned_org_node) "
                "VALUES ('t', 'p', 'g', 'goal', 'queued', 'planner')"
            )
            fresh.execute(
                "INSERT INTO work_graph_run (run_id, task_id, org_node, graph_manifest, status) "
                "VALUES ('r', 't', 'planner', '{}', 'planning')"
            )
            fresh.execute(
                "INSERT INTO artifact (artifact_id, artifact_type, task_id, produced_by_run_id, "
                "file_path, git_commit_sha) VALUES ('a', 'test_script', 't', 'r', 'x', 'sha')"
            )
            fresh.commit()
        finally:
            fresh.close()


class ExecutionLogFileUniqueConstraintTest(ConstraintTestBase):
    def setUp(self):
        super().setUp()
        self.insert_project()
        self.insert_task()
        self.insert_skill()
        self.insert_work_graph_run()
        self.conn.execute(
            "INSERT INTO work_node_execution "
            "(node_exec_id, run_id, node_id, subagent_type, skill_ref, "
            " input_schema_ref, output_schema_ref, input_payload, status) "
            "VALUES ('node-1', 'run-1', 'n1', 'coder', 'skill-1', 'in.json', 'out.json', '{}', 'running')"
        )
        self.conn.commit()

    def test_second_log_file_for_same_node_exec_rejected(self):
        self.conn.execute(
            "INSERT INTO execution_log_file (log_file_id, node_exec_id, project_id, file_path, format, status) "
            "VALUES ('log-1', 'node-1', 'proj-1', '/logs/node-1.jsonl', 'jsonl', 'active')"
        )
        self.conn.commit()
        with self.assertRaises(sqlite3.IntegrityError):
            self.conn.execute(
                "INSERT INTO execution_log_file (log_file_id, node_exec_id, project_id, file_path, format, status) "
                "VALUES ('log-2', 'node-1', 'proj-1', '/logs/node-1-again.jsonl', 'jsonl', 'active')"
            )


class BugRecordRootCauseCategoryTest(ConstraintTestBase):
    """core_engine was PH5-scoped and deliberately excluded pre-U-034
    (confirmed with user, 2026-07-30); U-034 (2026-08-02) widened the enum
    to accept it -- 05 §5.8's self-target routing needs it as a real,
    storable category, not just a documented intention."""

    def _insert_bug(self, root_cause_category):
        self.conn.execute(
            "INSERT INTO bug_record (bug_id, source, description, root_cause_category, status) "
            "VALUES ('bug-1', 'validator', 'wall collision misjudged', ?, 'open')",
            (root_cause_category,),
        )

    def test_all_four_categories_accepted(self):
        for category in ("planning", "coder", "skill_defect", "core_engine"):
            with self.subTest(category=category):
                conn = schema.connect(":memory:")
                schema.create_schema(conn)
                conn.execute(
                    "INSERT INTO bug_record (bug_id, source, description, root_cause_category, status) "
                    "VALUES ('bug-1', 'validator', 'x', ?, 'open')",
                    (category,),
                )
                conn.commit()
                conn.close()

    def test_unknown_category_still_rejected(self):
        with self.assertRaises(sqlite3.IntegrityError):
            self._insert_bug("not_a_real_category")


class DeploymentEnumConstraintTest(ConstraintTestBase):
    def _insert_deployment(self, deployment_id, status):
        self.conn.execute(
            "INSERT INTO deployment (deployment_id, git_commit_hash, previous_commit_hash, status, "
            "previous_instance_dir) VALUES (?, 'sha-new', 'sha-old', ?, '/tmp/prev-instance')",
            (deployment_id, status),
        )

    def test_every_state_machine_status_accepted(self):
        statuses = (
            "staged", "candidate_starting", "health_check_running", "health_check_passed",
            "health_check_failed", "cutover_in_progress", "live", "monitoring",
            "anomaly_detected", "rollback_in_progress", "rolled_back", "rollback_failed",
            "blocked_pending_human", "monitoring_window_clean", "finalized",
        )
        for status in statuses:
            with self.subTest(status=status):
                conn = schema.connect(":memory:")
                schema.create_schema(conn)
                conn.execute(
                    "INSERT INTO deployment (deployment_id, git_commit_hash, previous_commit_hash, status, "
                    "previous_instance_dir) VALUES ('d1', 'sha-new', 'sha-old', ?, '/tmp/prev-instance')",
                    (status,),
                )
                conn.commit()
                conn.close()

    def test_invalid_status_rejected(self):
        with self.assertRaises(sqlite3.IntegrityError):
            self._insert_deployment("d1", "not_a_real_status")


class DeploymentActiveUniqueIndexTest(ConstraintTestBase):
    """01-M 1.15: only one non-terminal deployment (not finalized/rolled_back/
    blocked_pending_human) at a time -- concurrent self-modifying deploys
    would corrupt the single Live/previous_instance_dir bookkeeping."""

    def _insert_deployment(self, deployment_id, status):
        self.conn.execute(
            "INSERT INTO deployment (deployment_id, git_commit_hash, previous_commit_hash, status, "
            "previous_instance_dir) VALUES (?, 'sha-new', 'sha-old', ?, '/tmp/prev-instance')",
            (deployment_id, status),
        )

    def test_second_active_deployment_rejected(self):
        self._insert_deployment("d1", "staged")
        self.conn.commit()
        with self.assertRaises(sqlite3.IntegrityError):
            self._insert_deployment("d2", "candidate_starting")

    def test_second_deployment_allowed_once_first_is_terminal(self):
        for terminal_status in ("finalized", "rolled_back", "blocked_pending_human"):
            with self.subTest(terminal_status=terminal_status):
                conn = schema.connect(":memory:")
                schema.create_schema(conn)
                conn.execute(
                    "INSERT INTO deployment (deployment_id, git_commit_hash, previous_commit_hash, status, "
                    "previous_instance_dir) VALUES ('d1', 'sha-new', 'sha-old', ?, '/tmp/prev-instance')",
                    (terminal_status,),
                )
                conn.commit()
                conn.execute(
                    "INSERT INTO deployment (deployment_id, git_commit_hash, previous_commit_hash, status, "
                    "previous_instance_dir) VALUES ('d2', 'sha-newer', 'sha-new', 'staged', '/tmp/prev-instance-2')"
                )
                conn.commit()
                count = conn.execute("SELECT COUNT(*) FROM deployment").fetchone()[0]
                conn.close()
                self.assertEqual(count, 2)


class DeployHealthSignalTest(ConstraintTestBase):
    def setUp(self):
        super().setUp()
        self.conn.execute(
            "INSERT INTO deployment (deployment_id, git_commit_hash, previous_commit_hash, status, "
            "previous_instance_dir) VALUES ('d1', 'sha-new', 'sha-old', 'health_check_running', "
            "'/tmp/prev-instance')"
        )
        self.conn.commit()

    def test_valid_signal_accepted(self):
        self.conn.execute(
            "INSERT INTO deploy_health_signal (signal_id, deployment_id, sampled_at, signal_type, status, detail) "
            "VALUES ('sig-1', 'd1', '2026-08-02T00:00:00Z', 'process_alive', 'ok', 'pid 1234 alive')"
        )
        self.conn.commit()
        row = self.conn.execute("SELECT status FROM deploy_health_signal WHERE signal_id = 'sig-1'").fetchone()
        self.assertEqual(row[0], "ok")

    def test_invalid_signal_type_rejected(self):
        with self.assertRaises(sqlite3.IntegrityError):
            self.conn.execute(
                "INSERT INTO deploy_health_signal (signal_id, deployment_id, sampled_at, signal_type, status) "
                "VALUES ('sig-1', 'd1', '2026-08-02T00:00:00Z', 'not_a_real_signal', 'ok')"
            )

    def test_invalid_status_rejected(self):
        with self.assertRaises(sqlite3.IntegrityError):
            self.conn.execute(
                "INSERT INTO deploy_health_signal (signal_id, deployment_id, sampled_at, signal_type, status) "
                "VALUES ('sig-1', 'd1', '2026-08-02T00:00:00Z', 'process_alive', 'not_a_real_status')"
            )

    def test_rejects_unknown_deployment_id(self):
        with self.assertRaises(sqlite3.IntegrityError):
            self.conn.execute(
                "INSERT INTO deploy_health_signal (signal_id, deployment_id, sampled_at, signal_type, status) "
                "VALUES ('sig-1', 'no-such-deployment', '2026-08-02T00:00:00Z', 'process_alive', 'ok')"
            )


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