vectorless_RAG / tests /test_client_isolation.py
amer003100's picture
feat(documents): add bulk DELETE /documents endpoint and active task cancellation
9f0929b
Raw History Blame Contribute Delete
7.42 kB
import uuid
import pytest
from fastapi.testclient import TestClient
from app import app
@pytest.fixture
def client():
return TestClient(app)
def test_public_endpoints_accessible_without_client_id(client):
"""Verifies that public endpoints (/, /health, /docs, /openapi.json) require no X-Client-ID."""
r_root = client.get("/")
assert r_root.status_code == 200
assert "message" in r_root.json()
r_health = client.get("/health")
assert r_health.status_code == 200
assert r_health.json()["status"] == "ok"
r_docs = client.get("/docs")
assert r_docs.status_code == 200
r_openapi = client.get("/openapi.json")
assert r_openapi.status_code == 200
def test_protected_endpoints_fallback_missing_client_id(client):
"""Protected endpoints gracefully fall back to default client ID when X-Client-ID is missing."""
r = client.get("/chats")
assert r.status_code == 200
assert "chats" in r.json()
r_docs = client.get("/documents")
assert r_docs.status_code == 200
assert "documents" in r_docs.json()
r_tools = client.post("/tools/calculate-math", json={"expression": "1 + 1"})
assert r_tools.status_code == 200
assert r_tools.json()["result"] == 2
def test_protected_endpoints_fallback_invalid_uuid(client):
"""Protected endpoints gracefully handle non-UUID string without throwing 400 error."""
bad_headers = {"X-Client-ID": "custom-client-string"}
r = client.get("/chats", headers=bad_headers)
assert r.status_code == 200
assert "chats" in r.json()
def test_chat_isolation_between_clients(client):
"""Client A and Client B must have complete data isolation for chats and messages."""
client_a_id = str(uuid.uuid4())
client_b_id = str(uuid.uuid4())
headers_a = {"X-Client-ID": client_a_id}
headers_b = {"X-Client-ID": client_b_id}
# 1. Client A creates a chat
create_res = client.post("/chat/new", headers=headers_a)
assert create_res.status_code == 200
chat_a_id = create_res.json()["chat_id"]
# 2. Client A lists chats -> finds chat_a_id
list_a = client.get("/chats", headers=headers_a).json()["chats"]
assert any(c["chat_id"] == chat_a_id for c in list_a)
# 3. Client B lists chats -> must NOT find chat_a_id
list_b = client.get("/chats", headers=headers_b).json()["chats"]
assert not any(c["chat_id"] == chat_a_id for c in list_b)
# 4. Client B attempts to fetch Client A's chat history -> 404
hist_b = client.get(f"/chat/{chat_a_id}/history", headers=headers_b)
assert hist_b.status_code == 404
# 5. Client B attempts to delete Client A's chat -> 404
del_res = client.delete(f"/chat/{chat_a_id}", headers=headers_b)
assert del_res.status_code == 404
# 6. Client A renames chat -> succeeds
rename_res = client.patch(f"/chat/{chat_a_id}/title", headers=headers_a, json={"title": "Client A Topic"})
assert rename_res.status_code == 200
# 7. Client B attempts to rename Client A's chat -> 404
rename_b = client.patch(f"/chat/{chat_a_id}/title", headers=headers_b, json={"title": "Hacked Title"})
assert rename_b.status_code == 404
# 8. Client A deletes chat -> succeeds
del_res_a = client.delete(f"/chat/{chat_a_id}", headers=headers_a)
assert del_res_a.status_code == 200
def test_tools_endpoint_accessible_with_valid_client_id(client):
"""Tools endpoint succeeds with valid client_id."""
client_id = str(uuid.uuid4())
headers = {"X-Client-ID": client_id}
res = client.post("/tools/calculate-math", headers=headers, json={"expression": "10 * 5"})
assert res.status_code == 200
assert res.json()["result"] == 50
def test_document_and_memory_isolation_between_clients(client):
"""Client A's documents, dossiers, and summaries must remain invisible to Client B."""
from core.dependencies import get_storage
storage = get_storage()
client_a = str(uuid.uuid4())
client_b = str(uuid.uuid4())
doc_a_id = f"doc_test_{uuid.uuid4().hex[:8]}"
# Seed document metadata for Client A directly into storage
storage.create_document(
doc_id=doc_a_id,
filename="client_a_contract.pdf",
content_hash="hash_a_12345",
file_path="/tmp/fake_contract.pdf",
owner_id=client_a,
)
storage.save_knowledge_dossier(
doc_id=doc_a_id,
filename="client_a_contract.pdf",
executive_summary="Confidential contract for Client A",
key_entities=["Client A Corp"],
key_findings=["Revenue target $10M"],
topics=["Confidential", "Finance"],
owner_id=client_a,
)
# Client A lists documents -> sees doc_a_id
headers_a = {"X-Client-ID": client_a}
docs_a = client.get("/documents", headers=headers_a).json()["documents"]
assert any(d["doc_id"] == doc_a_id for d in docs_a)
# Client B lists documents -> does NOT see doc_a_id
headers_b = {"X-Client-ID": client_b}
docs_b = client.get("/documents", headers=headers_b).json()["documents"]
assert not any(d["doc_id"] == doc_a_id for d in docs_b)
# Client B attempts to fetch Client A's document details -> 404
get_doc_b = client.get(f"/documents/{doc_a_id}", headers=headers_b)
assert get_doc_b.status_code == 404
# Client B attempts to fetch Client A's dossier -> 404
dossier_b = client.get(f"/memory/dossier/{doc_a_id}", headers=headers_b)
assert dossier_b.status_code == 404
# Client A fetches dossier -> 200
dossier_a = client.get(f"/memory/dossier/{doc_a_id}", headers=headers_a)
assert dossier_a.status_code == 200
assert dossier_a.json()["executive_summary"] == "Confidential contract for Client A"
# Client B attempts to delete Client A's document -> 404
del_b = client.delete(f"/documents/{doc_a_id}", headers=headers_b)
assert del_b.status_code == 404
# Client A deletes document -> 200
del_a = client.delete(f"/documents/{doc_a_id}", headers=headers_a)
assert del_a.status_code == 200
def test_delete_all_documents_bulk_endpoint(client):
"""Verifies that DELETE /documents deletes all documents belonging to the requesting client."""
from core.dependencies import get_storage
storage = get_storage()
client_c = str(uuid.uuid4())
client_d = str(uuid.uuid4())
headers_c = {"X-Client-ID": client_c}
headers_d = {"X-Client-ID": client_d}
doc_c1 = f"doc_c1_{uuid.uuid4().hex[:8]}"
doc_c2 = f"doc_c2_{uuid.uuid4().hex[:8]}"
doc_d = f"doc_d_{uuid.uuid4().hex[:8]}"
storage.create_document(doc_id=doc_c1, filename="doc1.pdf", owner_id=client_c)
storage.create_document(doc_id=doc_c2, filename="doc2.pdf", owner_id=client_c)
storage.create_document(doc_id=doc_d, filename="doc_d.pdf", owner_id=client_d)
# Client C deletes all
del_res = client.delete("/documents", headers=headers_c)
assert del_res.status_code == 200
data = del_res.json()
assert data["count"] >= 2
# Client C documents list is now empty of c1 and c2
docs_c = client.get("/documents", headers=headers_c).json()["documents"]
assert not any(d["doc_id"] in [doc_c1, doc_c2] for d in docs_c)
# Client D document is still intact
docs_d = client.get("/documents", headers=headers_d).json()["documents"]
assert any(d["doc_id"] == doc_d for d in docs_d)
# Cleanup Client D
client.delete(f"/documents/{doc_d}", headers=headers_d)