Skip to content

Commit 867d3d3

Browse files
committed
Add more unit tests for CosmosDB
1 parent ec12bca commit 867d3d3

File tree

3 files changed

+197
-10
lines changed

3 files changed

+197
-10
lines changed

app/backend/chat_history/cosmosdb.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -22,11 +22,11 @@
2222
@authenticated
2323
async def post_chat_history(auth_claims: Dict[str, Any]):
2424
if not current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED]:
25-
return jsonify({"error": "Chat history not enabled"}), 405
25+
return jsonify({"error": "Chat history not enabled"}), 400
2626

2727
container: ContainerProxy = current_app.config[CONFIG_COSMOS_HISTORY_CONTAINER]
2828
if not container:
29-
return jsonify({"error": "Chat history not enabled"}), 405
29+
return jsonify({"error": "Chat history not enabled"}), 400
3030

3131
entra_oid = auth_claims.get("oid")
3232
if not entra_oid:
@@ -52,11 +52,11 @@ async def post_chat_history(auth_claims: Dict[str, Any]):
5252
@authenticated
5353
async def get_chat_history(auth_claims: Dict[str, Any]):
5454
if not current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED]:
55-
return jsonify({"error": "Chat history not enabled"}), 405
55+
return jsonify({"error": "Chat history not enabled"}), 400
5656

5757
container: ContainerProxy = current_app.config[CONFIG_COSMOS_HISTORY_CONTAINER]
5858
if not container:
59-
return jsonify({"error": "Chat history not enabled"}), 405
59+
return jsonify({"error": "Chat history not enabled"}), 400
6060

6161
entra_oid = auth_claims.get("oid")
6262
if not entra_oid:
@@ -92,7 +92,7 @@ async def get_chat_history(auth_claims: Dict[str, Any]):
9292
}
9393
)
9494

95-
# If there are no page, StopAsyncIteration is raised
95+
# If there are no more pages, StopAsyncIteration is raised
9696
except StopAsyncIteration:
9797
items = []
9898
continuation_token = None
@@ -107,11 +107,11 @@ async def get_chat_history(auth_claims: Dict[str, Any]):
107107
@authenticated
108108
async def get_chat_history_session(auth_claims: Dict[str, Any], item_id: str):
109109
if not current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED]:
110-
return jsonify({"error": "Chat history not enabled"}), 405
110+
return jsonify({"error": "Chat history not enabled"}), 400
111111

112112
container: ContainerProxy = current_app.config[CONFIG_COSMOS_HISTORY_CONTAINER]
113113
if not container:
114-
return jsonify({"error": "Chat history not enabled"}), 405
114+
return jsonify({"error": "Chat history not enabled"}), 400
115115

116116
if not item_id:
117117
return jsonify({"error": "Invalid item ID specifier"}), 400
@@ -142,11 +142,11 @@ async def get_chat_history_session(auth_claims: Dict[str, Any], item_id: str):
142142
@authenticated
143143
async def delete_chat_history_session(auth_claims: Dict[str, Any], item_id: str):
144144
if not current_app.config[CONFIG_CHAT_HISTORY_COSMOS_ENABLED]:
145-
return jsonify({"error": "Chat history not enabled"}), 405
145+
return jsonify({"error": "Chat history not enabled"}), 400
146146

147147
container: ContainerProxy = current_app.config[CONFIG_COSMOS_HISTORY_CONTAINER]
148148
if not container:
149-
return jsonify({"error": "Chat history not enabled"}), 405
149+
return jsonify({"error": "Chat history not enabled"}), 400
150150

151151
if not item_id:
152152
return jsonify({"error": "Invalid path"}), 400
Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
{
2+
"continuation_token": "next",
3+
"items": [
4+
{
5+
"entra_oid": "OID_X",
6+
"id": "123",
7+
"timestamp": 123456789,
8+
"title": "This is a test message"
9+
}
10+
]
11+
}

tests/test_cosmosdb.py

Lines changed: 177 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,46 @@
1+
import json
2+
13
import pytest
24
from azure.cosmos.aio import ContainerProxy
35

6+
from .mocks import MockAsyncPageIterator
7+
8+
9+
class MockCosmosDBResultsIterator:
10+
def __init__(self):
11+
self.data = [
12+
[
13+
{
14+
"id": "123",
15+
"entra_oid": "OID_X",
16+
"title": "This is a test message",
17+
"timestamp": 123456789,
18+
"answers": [["This is a test message"]],
19+
}
20+
]
21+
]
22+
23+
def __aiter__(self):
24+
return self
25+
26+
async def __anext__(self):
27+
if not self.data:
28+
raise StopAsyncIteration
29+
return MockAsyncPageIterator(self.data.pop(0))
30+
31+
async def get_count(self):
32+
return len(self.data)
33+
34+
def by_page(self, continuation_token=None):
35+
if continuation_token:
36+
self.continuation_token = continuation_token + "next"
37+
else:
38+
self.continuation_token = "next"
39+
return self
40+
441

542
@pytest.mark.asyncio
6-
async def test_chathistory(auth_public_documents_client, monkeypatch):
43+
async def test_chathistory_newitem(auth_public_documents_client, monkeypatch):
744

845
async def mock_upsert_item(container_proxy, item, **kwargs):
946
assert item["id"] == "123"
@@ -22,3 +59,142 @@ async def mock_upsert_item(container_proxy, item, **kwargs):
2259
},
2360
)
2461
assert response.status_code == 201
62+
63+
64+
@pytest.mark.asyncio
65+
async def test_chathistory_newitem_error_disabled(client, monkeypatch):
66+
67+
response = await client.post(
68+
"/chat_history",
69+
headers={"Authorization": "Bearer MockToken"},
70+
json={
71+
"id": "123",
72+
"answers": [["This is a test message"]],
73+
},
74+
)
75+
assert response.status_code == 400
76+
77+
78+
@pytest.mark.asyncio
79+
async def test_chathistory_newitem_error_container(auth_public_documents_client, monkeypatch):
80+
auth_public_documents_client.app.config["cosmos_history_container"] = None
81+
response = await auth_public_documents_client.post(
82+
"/chat_history",
83+
headers={"Authorization": "Bearer MockToken"},
84+
json={
85+
"id": "123",
86+
"answers": [["This is a test message"]],
87+
},
88+
)
89+
assert response.status_code == 400
90+
91+
92+
@pytest.mark.asyncio
93+
async def test_chathistory_newitem_error_entra(auth_public_documents_client, monkeypatch):
94+
response = await auth_public_documents_client.post(
95+
"/chat_history",
96+
json={
97+
"id": "123",
98+
"answers": [["This is a test message"]],
99+
},
100+
)
101+
assert response.status_code == 401
102+
103+
104+
@pytest.mark.asyncio
105+
async def test_chathistory_newitem_error_runtime(auth_public_documents_client, monkeypatch):
106+
107+
async def mock_upsert_item(container_proxy, item, **kwargs):
108+
raise Exception("Test Exception")
109+
110+
monkeypatch.setattr(ContainerProxy, "upsert_item", mock_upsert_item)
111+
112+
response = await auth_public_documents_client.post(
113+
"/chat_history",
114+
headers={"Authorization": "Bearer MockToken"},
115+
json={
116+
"id": "123",
117+
"answers": [["This is a test message"]],
118+
},
119+
)
120+
assert response.status_code == 500
121+
assert (await response.get_json()) == {
122+
"error": "The app encountered an error processing your request.\nIf you are an administrator of the app, view the full error in the logs. See aka.ms/appservice-logs for more information.\nError type: <class 'Exception'>\n"
123+
}
124+
125+
126+
@pytest.mark.asyncio
127+
async def test_chathistory_query(auth_public_documents_client, monkeypatch, snapshot):
128+
129+
def mock_query_items(container_proxy, query, **kwargs):
130+
return MockCosmosDBResultsIterator()
131+
132+
monkeypatch.setattr(ContainerProxy, "query_items", mock_query_items)
133+
134+
response = await auth_public_documents_client.post(
135+
"/chat_history/items",
136+
headers={"Authorization": "Bearer MockToken"},
137+
json={"count": 20},
138+
)
139+
assert response.status_code == 200
140+
result = await response.get_json()
141+
snapshot.assert_match(json.dumps(result, indent=4), "result.json")
142+
143+
144+
@pytest.mark.asyncio
145+
async def test_chathistory_query_error_disabled(client, monkeypatch):
146+
147+
response = await client.post(
148+
"/chat_history/items",
149+
headers={"Authorization": "Bearer MockToken"},
150+
json={
151+
"id": "123",
152+
"answers": [["This is a test message"]],
153+
},
154+
)
155+
assert response.status_code == 400
156+
157+
158+
@pytest.mark.asyncio
159+
async def test_chathistory_query_error_container(auth_public_documents_client, monkeypatch):
160+
auth_public_documents_client.app.config["cosmos_history_container"] = None
161+
response = await auth_public_documents_client.post(
162+
"/chat_history/items",
163+
headers={"Authorization": "Bearer MockToken"},
164+
json={
165+
"id": "123",
166+
"answers": [["This is a test message"]],
167+
},
168+
)
169+
assert response.status_code == 400
170+
171+
172+
@pytest.mark.asyncio
173+
async def test_chathistory_query_error_entra(auth_public_documents_client, monkeypatch):
174+
response = await auth_public_documents_client.post(
175+
"/chat_history/items",
176+
json={
177+
"id": "123",
178+
"answers": [["This is a test message"]],
179+
},
180+
)
181+
assert response.status_code == 401
182+
183+
184+
@pytest.mark.asyncio
185+
async def test_chathistory_query_error_runtime(auth_public_documents_client, monkeypatch):
186+
187+
def mock_query_items(container_proxy, query, **kwargs):
188+
raise Exception("Test Exception")
189+
190+
monkeypatch.setattr(ContainerProxy, "query_items", mock_query_items)
191+
192+
response = await auth_public_documents_client.post(
193+
"/chat_history/items",
194+
headers={"Authorization": "Bearer MockToken"},
195+
json={"count": 20},
196+
)
197+
assert response.status_code == 500
198+
assert (await response.get_json()) == {
199+
"error": "The app encountered an error processing your request.\nIf you are an administrator of the app, view the full error in the logs. See aka.ms/appservice-logs for more information.\nError type: <class 'Exception'>\n"
200+
}

0 commit comments

Comments
 (0)