diff --git a/tests/conftest.py b/tests/conftest.py index 54c664d5..9a500609 100644 --- a/tests/conftest.py +++ b/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() diff --git a/tests/test_client.py b/tests/test_client.py index 9ba04416..f23a350e 100644 --- a/tests/test_client.py +++ b/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"), + ) + ]