import unittest
from unittest.mock import MagicMock, patch

import app.services.ops_agent_service as ops_agent_service
import app.services.ops_job_service as ops_job_service


class OpsSchemaInitTests(unittest.TestCase):
    @patch("app.services.ops_job_service.get_db")
    def test_ensure_ops_schema_uses_advisory_lock_and_skips_after_ready(self, mock_get_db) -> None:
        conn = MagicMock()
        cursor = MagicMock()
        db_ctx = MagicMock()
        cursor_ctx = MagicMock()
        db_ctx.__enter__.return_value = conn
        db_ctx.__exit__.return_value = False
        cursor_ctx.__enter__.return_value = cursor
        cursor_ctx.__exit__.return_value = False
        conn.cursor.return_value = cursor_ctx
        mock_get_db.return_value = db_ctx

        previous_ready = ops_job_service._OPS_SCHEMA_READY
        ops_job_service._OPS_SCHEMA_READY = False
        try:
            ops_job_service.ensure_ops_schema()
            ops_job_service.ensure_ops_schema()
        finally:
            ops_job_service._OPS_SCHEMA_READY = previous_ready

        self.assertEqual(1, mock_get_db.call_count)
        cursor.execute.assert_any_call(
            "SELECT pg_advisory_xact_lock(%s)",
            (ops_job_service._OPS_SCHEMA_ADVISORY_LOCK_KEY,),
        )
        cursor.execute.assert_any_call(ops_job_service._OPS_SCHEMA_SQL)
        conn.commit.assert_called_once()

    @patch("app.services.ops_agent_service.ensure_ops_schema")
    @patch("app.services.ops_agent_service.get_db")
    def test_ensure_ops_agent_schema_uses_advisory_lock_and_skips_after_ready(
        self,
        mock_get_db,
        mock_ensure_ops_schema,
    ) -> None:
        conn = MagicMock()
        cursor = MagicMock()
        db_ctx = MagicMock()
        cursor_ctx = MagicMock()
        db_ctx.__enter__.return_value = conn
        db_ctx.__exit__.return_value = False
        cursor_ctx.__enter__.return_value = cursor
        cursor_ctx.__exit__.return_value = False
        conn.cursor.return_value = cursor_ctx
        mock_get_db.return_value = db_ctx

        previous_ready = ops_agent_service._OPS_AGENT_SCHEMA_READY
        ops_agent_service._OPS_AGENT_SCHEMA_READY = False
        try:
            ops_agent_service.ensure_ops_agent_schema()
            ops_agent_service.ensure_ops_agent_schema()
        finally:
            ops_agent_service._OPS_AGENT_SCHEMA_READY = previous_ready

        self.assertEqual(2, mock_ensure_ops_schema.call_count)
        self.assertEqual(1, mock_get_db.call_count)
        cursor.execute.assert_any_call(
            "SELECT pg_advisory_xact_lock(%s)",
            (ops_agent_service._OPS_AGENT_SCHEMA_ADVISORY_LOCK_KEY,),
        )
        cursor.execute.assert_any_call(ops_agent_service._AGENT_SCHEMA_SQL)
        conn.commit.assert_called_once()


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