From 3678aa5614858ac75189a9e0f5518c3138072f56 Mon Sep 17 00:00:00 2001 From: sabaimran Date: Sat, 4 Nov 2023 14:29:30 -0700 Subject: [PATCH] Add tests to validate expected behaviors in the multi-user scenario --- tests/conftest.py | 13 ++++ tests/test_client.py | 33 +++++------ tests/test_multiple_users.py | 111 +++++++++++++++++++++++++++++++++++ 3 files changed, 140 insertions(+), 17 deletions(-) create mode 100644 tests/test_multiple_users.py diff --git a/tests/conftest.py b/tests/conftest.py index a5f23dd2..3c579834 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -98,6 +98,19 @@ def api_user(default_user): ) +@pytest.mark.django_db +@pytest.fixture +def api_user2(default_user2): + if KhojApiUser.objects.filter(user=default_user2).exists(): + return KhojApiUser.objects.get(user=default_user2) + + return KhojApiUser.objects.create( + user=default_user2, + name="api-key", + token="kk-diff-secret", + ) + + @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 5cf438c7..c105c605 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -352,21 +352,20 @@ def test_different_user_data_not_accessed(client, sample_org_data, default_user: def get_sample_files_data(): - return { - "files": ("path/to/filename.org", "* practicing piano", "text/org"), - "files": ("path/to/filename1.org", "** top 3 reasons why I moved to SF", "text/org"), - "files": ("path/to/filename2.org", "* how to build a search engine", "text/org"), - "files": ("path/to/filename.pdf", "Moore's law does not apply to consumer hardware", "application/pdf"), - "files": ("path/to/filename1.pdf", "The sun is a ball of helium", "application/pdf"), - "files": ("path/to/filename2.pdf", "Effect of sunshine on baseline human happiness", "application/pdf"), - "files": ("path/to/filename.txt", "data,column,value", "text/plain"), - "files": ("path/to/filename1.txt", "my first web page", "text/plain"), - "files": ("path/to/filename2.txt", "2021-02-02 Journal Entry", "text/plain"), - "files": ("path/to/filename.md", "# Notes from client call", "text/markdown"), - "files": ( - "path/to/filename1.md", - "## Studying anthropological records from the Fatimid caliphate", - "text/markdown", + return [ + ("files", ("path/to/filename.org", "* practicing piano", "text/org")), + ("files", ("path/to/filename1.org", "** top 3 reasons why I moved to SF", "text/org")), + ("files", ("path/to/filename2.org", "* how to build a search engine", "text/org")), + ("files", ("path/to/filename.pdf", "Moore's law does not apply to consumer hardware", "application/pdf")), + ("files", ("path/to/filename1.pdf", "The sun is a ball of helium", "application/pdf")), + ("files", ("path/to/filename2.pdf", "Effect of sunshine on baseline human happiness", "application/pdf")), + ("files", ("path/to/filename.txt", "data,column,value", "text/plain")), + ("files", ("path/to/filename1.txt", "my first web page", "text/plain")), + ("files", ("path/to/filename2.txt", "2021-02-02 Journal Entry", "text/plain")), + ("files", ("path/to/filename.md", "# Notes from client call", "text/markdown")), + ( + "files", + ("path/to/filename1.md", "## Studying anthropological records from the Fatimid caliphate", "text/markdown"), ), - "files": ("path/to/filename2.md", "**Understanding science through the lens of art**", "text/markdown"), - } + ("files", ("path/to/filename2.md", "**Understanding science through the lens of art**", "text/markdown")), + ] diff --git a/tests/test_multiple_users.py b/tests/test_multiple_users.py new file mode 100644 index 00000000..95a2535f --- /dev/null +++ b/tests/test_multiple_users.py @@ -0,0 +1,111 @@ +# Standard Modules +from io import BytesIO +from PIL import Image +from urllib.parse import quote +import pytest + +# External Packages +from fastapi.testclient import TestClient +from fastapi import FastAPI, UploadFile +from io import BytesIO +import pytest + +# Internal Packages +from khoj.configure import configure_routes, configure_search_types +from khoj.utils import state +from khoj.utils.state import search_models, content_index, config +from khoj.search_type import text_search, image_search +from khoj.utils.rawconfig import ContentConfig, SearchConfig +from khoj.processor.org_mode.org_to_entries import OrgToEntries +from database.models import KhojUser, KhojApiUser +from database.adapters import EntryAdapters + + +# ---------------------------------------------------------------------------------------------------- +@pytest.mark.django_db(transaction=True) +def test_search_for_user2_returns_empty(client, api_user2: KhojApiUser): + token = api_user2.token + headers = {"Authorization": f"Bearer {token}"} + for content_type in ["all", "org", "markdown", "pdf", "github", "notion", "plaintext"]: + # Act + response = client.get(f"/api/search?q=random&t={content_type}", headers=headers) + # Assert + assert response.text == "[]" + assert response.status_code == 200, f"Returned status: {response.status_code} for content type: {content_type}" + + +# ---------------------------------------------------------------------------------------------------- +@pytest.mark.django_db(transaction=True) +def test_index_update_with_user2(client, api_user2: KhojApiUser): + # Arrange + files = get_sample_files_data() + source_file_symbol = set([f[1][0] for f in files]) + + headers = {"Authorization": f"Bearer {api_user2.token}"} + update_response = client.post("/api/v1/index/update", files=files, headers=headers) + search_response = client.get("/api/search?q=hardware&t=all", headers=headers) + results = search_response.json() + + # Assert + assert update_response.status_code == 200 + assert len(results) == 5 + for result in results: + assert result["additional"]["file"] in source_file_symbol + + +@pytest.mark.django_db(transaction=True) +def test_index_update_with_user2_inaccessible_user1(client, api_user2: KhojApiUser, api_user: KhojApiUser): + # Arrange + files = get_sample_files_data() + source_file_symbol = set([f[1][0] for f in files]) + + headers = {"Authorization": f"Bearer {api_user2.token}"} + update_response = client.post("/api/v1/index/update", files=files, headers=headers) + + # Act + headers = {"Authorization": f"Bearer {api_user.token}"} + search_response = client.get("/api/search?q=hardware&t=all", headers=headers) + results = search_response.json() + + # Assert + assert update_response.status_code == 200 + assert len(results) == 4 + for result in results: + assert result["additional"]["file"] not in source_file_symbol + + +# ---------------------------------------------------------------------------------------------------- +@pytest.mark.django_db(transaction=True) +def test_different_user_data_not_accessed(client, sample_org_data, default_user: KhojUser): + # Arrange + headers = {"Authorization": "Bearer kk-token"} # Token for default_user2 + text_search.setup(OrgToEntries, sample_org_data, regenerate=False, user=default_user) + user_query = quote("How to git install application?") + + # Act + response = client.get(f"/api/search?q={user_query}&n=1&t=org", headers=headers) + + # Assert + assert response.status_code == 403 + # assert actual response has no data as the default_user is different from the user making the query (anonymous) + assert len(response.json()) == 1 and response.json()["detail"] == "Forbidden" + + +def get_sample_files_data(): + return [ + ("files", ("path/to/filename.org", "* practicing piano", "text/org")), + ("files", ("path/to/filename1.org", "** top 3 reasons why I moved to SF", "text/org")), + ("files", ("path/to/filename2.org", "* how to build a search engine", "text/org")), + ("files", ("path/to/filename.pdf", "Moore's law does not apply to consumer hardware", "application/pdf")), + ("files", ("path/to/filename1.pdf", "The sun is a ball of helium", "application/pdf")), + ("files", ("path/to/filename2.pdf", "Effect of sunshine on baseline human happiness", "application/pdf")), + ("files", ("path/to/filename.txt", "data,column,value", "text/plain")), + ("files", ("path/to/filename1.txt", "my first web page", "text/plain")), + ("files", ("path/to/filename2.txt", "2021-02-02 Journal Entry", "text/plain")), + ("files", ("path/to/filename.md", "# Notes from client call", "text/markdown")), + ( + "files", + ("path/to/filename1.md", "## Studying anthropological records from the Fatimid caliphate", "text/markdown"), + ), + ("files", ("path/to/filename2.md", "**Understanding science through the lens of art**", "text/markdown")), + ]