diff --git a/tests/test_distributed_lock.py b/tests/test_distributed_lock.py new file mode 100644 index 00000000..475fc295 --- /dev/null +++ b/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()