Add additional unit tests to verify behavior of unsubscribed/subscribed users

tests/conftest.py

@@ -102,6 +102,24 @@ def default_user3():
     return user
 
 
+@pytest.mark.django_db
+@pytest.fixture
+def default_user4():
+    """
+    This user should not have a valid subscription
+    """
+    if KhojUser.objects.filter(username="default4").exists():
+        return KhojUser.objects.get(username="default4")
+
+    user = KhojUser.objects.create(
+        username="default4",
+        email="default4@example.com",
+        password="default4",
+    )
+    SubscriptionFactory(user=user, renewal_date=None)
+    return user
+
+
 @pytest.mark.django_db
 @pytest.fixture
 def api_user(default_user):
@@ -141,6 +159,19 @@ def api_user3(default_user3):
     )
 
 
+@pytest.mark.django_db
+@pytest.fixture
+def api_user4(default_user4):
+    if KhojApiUser.objects.filter(user=default_user4).exists():
+        return KhojApiUser.objects.get(user=default_user4)
+
+    return KhojApiUser.objects.create(
+        user=default_user4,
+        name="api-key",
+        token="kk-diff-secret-4",
+    )
+
+
 @pytest.fixture(scope="session")
 def search_models(search_config: SearchConfig):
     search_models = SearchModels()

tests/test_client.py

@@ -140,6 +140,38 @@ def test_index_update_big_files(client):
     assert response.status_code == 429
 
 
+# ----------------------------------------------------------------------------------------------------
+@pytest.mark.django_db(transaction=True)
+def test_index_update_medium_file_unsubscribed(client, api_user4: KhojApiUser):
+    # Arrange
+    api_token = api_user4.token
+    state.billing_enabled = True
+    files = get_medium_size_sample_files_data()
+    headers = {"Authorization": f"Bearer {api_token}"}
+
+    # Act
+    response = client.post("/api/v1/index/update", files=files, headers=headers)
+
+    # Assert
+    assert response.status_code == 429
+
+
+# ----------------------------------------------------------------------------------------------------
+@pytest.mark.django_db(transaction=True)
+def test_index_update_normal_file_unsubscribed(client, api_user4: KhojApiUser):
+    # Arrange
+    api_token = api_user4.token
+    state.billing_enabled = True
+    files = get_sample_files_data()
+    headers = {"Authorization": f"Bearer {api_token}"}
+
+    # Act
+    response = client.post("/api/v1/index/update", files=files, headers=headers)
+
+    # Assert
+    assert response.status_code == 200
+
+
 @pytest.mark.django_db(transaction=True)
 def test_index_update_big_files_no_billing(client):
     # Arrange
@@ -460,3 +492,13 @@ def get_big_size_sample_files_data():
             ("path/to/filename.org", big_text, "text/org"),
         )
     ]
+
+
+def get_medium_size_sample_files_data():
+    big_text = "a" * (10 * 1024 * 1024)  # a string of approximately 10 MB
+    return [
+        (
+            "files",
+            ("path/to/filename.org", big_text, "text/org"),
+        )
+    ]