Add tests for the db lock

tests/conftest.py

@@ -35,6 +35,7 @@ from tests.helpers import (
     ChatModelOptionsFactory,
     OfflineChatProcessorConversationConfigFactory,
     OpenAIProcessorConversationConfigFactory,
+    ProcessLockFactory,
     SubscriptionFactory,
     UserConversationProcessorConfigFactory,
     UserFactory,
@@ -206,6 +207,12 @@ def search_models(search_config: SearchConfig):
     return search_models
 
 
+@pytest.mark.django_db
+@pytest.fixture
+def default_process_lock():
+    return ProcessLockFactory()
+
+
 @pytest.fixture
 def anyio_backend():
     return "asyncio"

tests/helpers.py

@@ -11,6 +11,7 @@ from khoj.database.models import (
     KhojUser,
     OfflineChatProcessorConversationConfig,
     OpenAIProcessorConversationConfig,
+    ProcessLock,
     SearchModelConfig,
     Subscription,
     UserConversationConfig,
@@ -93,3 +94,10 @@ class SubscriptionFactory(factory.django.DjangoModelFactory):
     type = "standard"
     is_recurring = False
     renewal_date = make_aware(datetime.strptime("2100-04-01", "%Y-%m-%d"))
+
+
+class ProcessLockFactory(factory.django.DjangoModelFactory):
+    class Meta:
+        model = ProcessLock
+
+    name = "test_lock"

tests/test_db_lock.py

@@ -0,0 +1,58 @@
+import time
+
+import pytest
+
+from khoj.database.adapters import ProcessLockAdapters
+from khoj.database.models import ProcessLock
+from tests.helpers import ProcessLockFactory
+
+
+@pytest.mark.django_db(transaction=True)
+def test_process_lock(default_process_lock):
+    # Arrange
+    lock: ProcessLock = default_process_lock
+
+    # Assert
+    assert True == ProcessLockAdapters.is_process_locked(lock.name)
+
+
+@pytest.mark.django_db(transaction=True)
+def test_expired_process_lock():
+    # Arrange
+    lock: ProcessLock = ProcessLockFactory(name="test_expired_lock", max_duration_in_seconds=2)
+
+    # Act
+    time.sleep(3)
+
+    # Assert
+    assert False == ProcessLockAdapters.is_process_locked(lock.name)
+
+
+@pytest.mark.django_db(transaction=True)
+def test_in_progress_lock(default_process_lock):
+    # Arrange
+    lock: ProcessLock = default_process_lock
+
+    # Act
+    ProcessLockAdapters.run_with_lock(lock.name, lambda: time.sleep(2))
+
+    # Assert
+    assert True == ProcessLockAdapters.is_process_locked(lock.name)
+
+
+@pytest.mark.django_db(transaction=True)
+def test_run_with_completed():
+    # Arrange
+    ProcessLockAdapters.run_with_lock("test_run_with", lambda: time.sleep(2))
+
+    # Act
+    time.sleep(4)
+
+    # Assert
+    assert False == ProcessLockAdapters.is_process_locked("test_run_with")
+
+
+@pytest.mark.django_db(transaction=True)
+def test_nonexistent_lock():
+    # Assert
+    assert False == ProcessLockAdapters.is_process_locked("nonexistent_lock")