test: cover startup distributed lock behavior

Author: tarjan1

tests/test_distributed_lock.py

@@ -0,0 +1,170 @@
+"""Focused tests for the startup embedding distributed lock."""
+
+import sys
+import unittest
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+from sqlalchemy.engine import Connection
+from sqlalchemy.exc import SQLAlchemyError
+
+
+BACKEND_DIR = Path(__file__).resolve().parents[1] / "backend"
+sys.path.insert(0, str(BACKEND_DIR))
+
+from common.utils import distributed_lock as lock_module  # noqa: E402
+from common.utils.distributed_lock import (  # noqa: E402
+    DistributedLock,
+    LockStatus,
+    SingleWorkerGuard,
+)
+
+
+class DistributedLockTestCase(unittest.TestCase):
+    """Verify PostgreSQL advisory-lock acquisition and cleanup behavior."""
+
+    @staticmethod
+    def _connection(acquired: bool = True) -> MagicMock:
+        connection = MagicMock(spec=Connection)
+        connection.closed = False
+        connection.execute.return_value.scalar_one.return_value = acquired
+        return connection
+
+    def test_advisory_key_is_stable_across_processes(self) -> None:
+        key = DistributedLock._postgres_advisory_key("sqlbot:startup:embedding")
+
+        self.assertEqual(6735965423478359195, key)
+        self.assertGreaterEqual(key, -(2**63))
+        self.assertLessEqual(key, 2**63 - 1)
+
+    def test_try_acquire_keeps_successful_connection_open(self) -> None:
+        connection = self._connection(acquired=True)
+
+        with patch.object(lock_module, "engine") as engine:
+            engine.connect.return_value = connection
+
+            acquired_lock = DistributedLock.try_acquire("startup-task")
+
+        self.assertIsInstance(acquired_lock, DistributedLock)
+        connection.commit.assert_called_once_with()
+        connection.close.assert_not_called()
+
+    def test_try_acquire_closes_connection_when_lock_is_busy(self) -> None:
+        connection = self._connection(acquired=False)
+
+        with patch.object(lock_module, "engine") as engine:
+            engine.connect.return_value = connection
+
+            acquired_lock = DistributedLock.try_acquire("startup-task")
+
+        self.assertIsNone(acquired_lock)
+        connection.commit.assert_called_once_with()
+        connection.close.assert_called_once_with()
+
+    def test_try_acquire_closes_connection_on_sqlalchemy_error(self) -> None:
+        connection = self._connection()
+        connection.execute.side_effect = SQLAlchemyError("database error")
+
+        with (
+            patch.object(lock_module, "engine") as engine,
+            patch.object(lock_module.SQLBotLogUtil, "exception") as log_error,
+        ):
+            engine.connect.return_value = connection
+
+            acquired_lock = DistributedLock.try_acquire("startup-task")
+
+        self.assertIsNone(acquired_lock)
+        connection.close.assert_called_once_with()
+        log_error.assert_called_once()
+
+    def test_release_unlocks_commits_and_closes_connection(self) -> None:
+        connection = self._connection()
+        acquired_lock = DistributedLock(1234, connection)
+
+        acquired_lock.release()
+
+        connection.execute.assert_called_once()
+        connection.commit.assert_called_once_with()
+        connection.close.assert_called_once_with()
+
+    def test_release_still_closes_connection_when_unlock_fails(self) -> None:
+        connection = self._connection()
+        connection.execute.side_effect = SQLAlchemyError("unlock error")
+        acquired_lock = DistributedLock(1234, connection)
+
+        with self.assertRaises(SQLAlchemyError):
+            acquired_lock.release()
+
+        connection.commit.assert_not_called()
+        connection.close.assert_called_once_with()
+
+
+class SingleWorkerGuardTestCase(unittest.TestCase):
+    """Verify that only the process holding the lock runs startup tasks."""
+
+    def setUp(self) -> None:
+        SingleWorkerGuard._lock_status = LockStatus.NOT_ATTEMPTED
+        SingleWorkerGuard._acquired_lock = None
+
+    def tearDown(self) -> None:
+        SingleWorkerGuard._lock_status = LockStatus.NOT_ATTEMPTED
+        SingleWorkerGuard._acquired_lock = None
+
+    def test_acquired_worker_runs_all_tasks_with_one_lock_attempt(self) -> None:
+        acquired_lock = MagicMock(spec=DistributedLock)
+        executed_tasks: list[str] = []
+
+        @SingleWorkerGuard.once
+        def first_task() -> None:
+            executed_tasks.append("first")
+
+        @SingleWorkerGuard.once
+        def second_task() -> None:
+            executed_tasks.append("second")
+
+        with patch.object(
+            DistributedLock,
+            "try_acquire",
+            return_value=acquired_lock,
+        ) as try_acquire:
+            first_task()
+            second_task()
+
+        self.assertEqual(["first", "second"], executed_tasks)
+        try_acquire.assert_called_once_with(SingleWorkerGuard._LOCK_KEY)
+
+    def test_non_acquired_worker_skips_tasks_without_retrying(self) -> None:
+        task = MagicMock()
+        guarded_task = SingleWorkerGuard.once(task)
+
+        with (
+            patch.object(
+                DistributedLock,
+                "try_acquire",
+                return_value=None,
+            ) as try_acquire,
+            patch.object(lock_module.SQLBotLogUtil, "info"),
+        ):
+            guarded_task()
+            guarded_task()
+
+        task.assert_not_called()
+        try_acquire.assert_called_once_with(SingleWorkerGuard._LOCK_KEY)
+
+    def test_release_releases_lock_and_resets_process_state(self) -> None:
+        acquired_lock = MagicMock(spec=DistributedLock)
+        SingleWorkerGuard._lock_status = LockStatus.ACQUIRED
+        SingleWorkerGuard._acquired_lock = acquired_lock
+
+        SingleWorkerGuard.release()
+
+        acquired_lock.release.assert_called_once_with()
+        self.assertIsNone(SingleWorkerGuard._acquired_lock)
+        self.assertIs(
+            LockStatus.NOT_ATTEMPTED,
+            SingleWorkerGuard._lock_status,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()