From 3132430737f05290ba13bec4849b251356b9a6d1 Mon Sep 17 00:00:00 2001
From: sabaimran <narmiabas@gmail.com>
Date: Wed, 17 Apr 2024 13:22:41 +0530
Subject: [PATCH] Add tests for the db lock

---
 tests/conftest.py     |  7 ++++++
 tests/helpers.py      |  8 ++++++
 tests/test_db_lock.py | 58 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+)
 create mode 100644 tests/test_db_lock.py

diff --git a/tests/conftest.py b/tests/conftest.py
index d513a574..2bf93928 100644
--- a/tests/conftest.py
+++ b/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"
diff --git a/tests/helpers.py b/tests/helpers.py
index 26c7e2af..0dbac55f 100644
--- a/tests/helpers.py
+++ b/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"
diff --git a/tests/test_db_lock.py b/tests/test_db_lock.py
new file mode 100644
index 00000000..3a5ff0b0
--- /dev/null
+++ b/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")