76 lines
2.8 KiB
Python
76 lines
2.8 KiB
Python
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()
|