File size: 13,020 Bytes
365de9c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 |
import os
from uuid import uuid4
import re
import httpx
import pytest
import pytest_asyncio
from app.backend.models.users import User # noqa: F401
from app.backend.models.chats import Chat # noqa: F401
from app.backend.models.messages import Message, get_messages_by_chat_id # noqa: F401
from app.settings import BASE_DIR
BASE_URL = os.environ.get('HF1_URL')
# --- Async Fixtures for Setup ---
@pytest_asyncio.fixture
async def artificial_user():
"""Fixture to create and log in an artificial user, returning user data."""
email = f"Test{uuid4()}@test.com"
password = "Goida123!"
payload = {"email": email, "password": password}
async with httpx.AsyncClient(verify=False) as client:
# Create user
response = await client.post(url=BASE_URL + "/new_user", json=payload, timeout=30.0)
if response.status_code != 200:
raise RuntimeError(f"Failed to create artificial user: {response.status_code} - {response.text}")
id = response.json().get("id", -1)
# Log in to get access token
headers = {"Content-Type": "application/json", "Accept": "application/json"}
response = await client.post(url=BASE_URL + "/login", json=payload, headers=headers, timeout=30.0)
if response.status_code != 200:
raise RuntimeError(f"Login failed: {response.status_code} - {response.text}")
# Extract access token from set-cookie header
set_cookie_header = response.headers.get("set-cookie")
access_token = None
if set_cookie_header and "access_token=" in set_cookie_header:
access_token = set_cookie_header.split("access_token=")[1].split(";")[0]
if not access_token:
raise RuntimeError("No access token received from login")
return {"id": id, "email": email, "password": password, "access_token": access_token}
@pytest_asyncio.fixture
async def chat_data(artificial_user):
"""Fixture to create a chat for the artificial user, returning chat data."""
cookie = {"access_token": artificial_user["access_token"]}
async with httpx.AsyncClient(verify=False, cookies=cookie) as client:
# Create chat
response = await client.post(url=BASE_URL + "/new_chat", timeout=30.0)
if response.status_code != 303:
raise RuntimeError(f"Error while trying to create chat: {response.status_code} - {response.text}")
redirect_to = response.headers.get("location")
if not redirect_to or "login" in redirect_to:
raise RuntimeError(f"Authentication failed, redirected to: {redirect_to}")
# Follow redirect
response = await client.get(url=BASE_URL + redirect_to)
if response.status_code != 200:
raise RuntimeError(f"Error while accessing chat: {response.status_code} - {response.text}")
try:
chat_id = int(redirect_to.split("/")[-1].split("id=")[-1])
except ValueError as e:
raise RuntimeError(f"Failed to parse chat_id from URL: {redirect_to} - {e}")
res = {"chat_id": chat_id, "user": artificial_user, "cookie": cookie}
print(res)
return res
# --- Async Test Functions ---
@pytest.mark.asyncio
async def test_create_artificial_user(artificial_user):
"""Test that an artificial user can be created and logged in."""
assert artificial_user["email"] is not None
assert artificial_user["password"] == "Goida123!"
assert artificial_user["access_token"] is not None
@pytest.mark.asyncio
async def test_validate_chat_creation(chat_data):
"""Test that a chat can be created successfully."""
assert chat_data["chat_id"] > 0
assert "access_token" in chat_data["cookie"]
@pytest.mark.asyncio
async def test_validate_message_sending(chat_data):
"""Test that a message can be sent to the chat."""
async with httpx.AsyncClient(verify=False, cookies=chat_data["cookie"]) as client:
payload = {"prompt": "How is your day?", "chat_id": chat_data["chat_id"]}
response = await client.post(
url=BASE_URL + "/message_with_docs",
data=payload,
timeout=180,
)
assert response.status_code == 200, f"Failed to send message: {response.status_code} - {response.text}"
@pytest.mark.asyncio
async def test_validate_docs_uploading(chat_data):
"""Test that a document can be uploaded with a message."""
file_path = os.path.join(BASE_DIR, "app", "tests", "integration", "testfile.txt")
# Create a test file if it doesn't exist
if not os.path.exists(file_path):
with open(file_path, "w") as f:
f.write("This is a test file for validation.")
async with httpx.AsyncClient(verify=False, cookies=chat_data["cookie"]) as client:
with open(file_path, "rb") as f:
form_fields = {
"prompt": "How is your day?",
"chat_id": str(chat_data["chat_id"]),
}
files = [("files", ("testfile.txt", f, "text/plain"))]
response = await client.post(
url=BASE_URL + "/message_with_docs",
data=form_fields,
files=files,
timeout=180,
)
assert response.status_code == 200, f"Failed to upload docs: {response.status_code} - {response.text}"
@pytest.mark.asyncio
async def test_validate_message_registration(chat_data):
"""Test that a sent message is registered in the chat."""
async with httpx.AsyncClient(verify=False, cookies=chat_data["cookie"]) as client:
initial = get_messages_by_chat_id(chat_data["chat_id"]).count()
payload = {"prompt": "How is your day?", "chat_id": chat_data["chat_id"]}
response = await client.post(
url=BASE_URL + "/message_with_docs",
data=payload,
timeout=180,
)
assert response.status_code == 200, f"Failed to send message: {response.status_code} - {response.text}"
after_sending = get_messages_by_chat_id(chat_data["chat_id"]).count()
assert after_sending - initial == 1, "Message was not registered"
@pytest.mark.asyncio
async def test_document_creation(chat_data):
"""Test that local storage is created properly and document is saved right to it."""
file_path = os.path.join(BASE_DIR, "app", "tests", "integration", "testfile.txt")
# Create a test file if it doesn't exist
if not os.path.exists(file_path):
with open(file_path, "w") as f:
f.write("This is a test file for validation.")
path_to_storage = os.path.join(BASE_DIR, "chats_storage", f"user_id={chat_data['user']['id']}", f"chat_id={chat_data['chat_id']}", "documents")
before_request = len(os.listdir(path_to_storage))
async with httpx.AsyncClient(verify=False, cookies=chat_data["cookie"]) as client:
with open(file_path, "rb") as f:
form_fields = {
"prompt": "How is your day?",
"chat_id": str(chat_data["chat_id"]),
}
files = [("files", ("testfile.txt", f, "text/plain"))]
response = await client.post(
url=BASE_URL + "/message_with_docs",
data=form_fields,
files=files,
timeout=180,
)
assert response.status_code == 200, f"Failed to upload docs: {response.status_code} - {response.text}"
after_request = len(os.listdir(path_to_storage))
assert after_request - before_request == 2
@pytest.mark.asyncio
async def test_document_content(chat_data):
"""Test that document that was sent is the same with the stored one."""
file_path = os.path.join(BASE_DIR, "app", "tests", "integration", "testfile.txt")
# Create a test file if it doesn't exist
if not os.path.exists(file_path):
with open(file_path, "w") as f:
f.write("This is a test file for validation.")
path_to_storage = os.path.join(BASE_DIR, "chats_storage", f"user_id={chat_data['user']['id']}", f"chat_id={chat_data['chat_id']}", "documents")
before_request = len(os.listdir(path_to_storage))
async with httpx.AsyncClient(verify=False, cookies=chat_data["cookie"]) as client:
with open(file_path, "rb") as f:
form_fields = {
"prompt": "How is your day?",
"chat_id": str(chat_data["chat_id"]),
}
files = [("files", ("testfile.txt", f, "text/plain"))]
response = await client.post(
url=BASE_URL + "/message_with_docs",
data=form_fields,
files=files,
timeout=180,
)
assert response.status_code == 200, f"Failed to upload docs: {response.status_code} - {response.text}"
after_request = len(os.listdir(path_to_storage))
assert after_request - before_request == 2
file_name = [file for file in os.listdir(path_to_storage) if "pdfs" not in file][0]
with open(os.path.join(path_to_storage, file_name), "r") as stored:
stored_content = stored.read()
with open(file_path, "r") as stored:
posted_content = stored.read()
assert posted_content == stored_content
@pytest.mark.asyncio
async def test_document_uploading_speed(chat_data):
"""Test that document uploading meets timeout requirements."""
file_path = os.path.join(BASE_DIR, "app", "tests", "integration", "testfile2.txt")
# Create a test file if it doesn't exist
if not os.path.exists(file_path):
with open(file_path, "w") as f:
f.write("This is a test file for validation.\n" * 1000000)
async with httpx.AsyncClient(verify=False, cookies=chat_data["cookie"]) as client:
with open(file_path, "rb") as f:
form_fields = {
"prompt": "How is your day?",
"chat_id": str(chat_data["chat_id"]),
}
files = [("files", ("testfile2.txt", f, "text/plain"))]
try:
response = await client.post(
url=BASE_URL + "/message_with_docs",
data=form_fields,
files=files,
timeout=50,
)
assert response.status_code == 200, f"Failed to upload docs: {response.status_code} - {response.text}"
except httpx.TimeoutException:
raise RuntimeError("The loading speed of large documents is too slow")
@pytest.mark.asyncio
async def test_xss_protection(chat_data):
"""Test that messages are protected from XSS attacks."""
async with httpx.AsyncClient(verify=False, cookies=chat_data["cookie"]) as client:
payload = {"prompt": "<script>How is your day?<script>", "chat_id": chat_data["chat_id"]}
response = await client.post(
url=BASE_URL + "/message_with_docs",
data=payload,
timeout=180,
)
assert response.status_code == 200, f"Failed to send message: {response.status_code} - {response.text}"
async with httpx.AsyncClient(verify=False, cookies=chat_data["cookie"]) as client:
payload = {"chat_id": chat_data["chat_id"]}
response = await client.post(
url=BASE_URL + f"/chats/id={chat_data['chat_id']}/history",
data=payload,
timeout=180,
)
response_json = response.json()
for message in response_json.get("history", [None]):
if message and "script" in message.get("content", ""):
raise RuntimeError("Messages are not protected from XSS attacks")
@pytest.mark.asyncio
async def test_source_citation(chat_data):
"""Test that a document can be uploaded with a message and cited."""
file_path = os.path.join(BASE_DIR, "app", "tests", "integration", "testfile3.txt")
# Create a test file if it doesn't exist
if os.path.exists(file_path):
with open(file_path, "w") as f:
f.write("How is your day: My day is good.")
async with httpx.AsyncClient(verify=False, cookies=chat_data["cookie"]) as client:
with open(file_path, "rb") as f:
form_fields = {
"prompt": "How is your day? CITE THE DOCUMENT THAT WILL BE ATTACHED!",
"chat_id": str(chat_data["chat_id"]),
}
files = [("files", ("testfile3.txt", f, "text/plain"))]
response = await client.post(
url=BASE_URL + "/message_with_docs",
data=form_fields,
files=files,
timeout=180,
)
assert response.status_code == 200, f"Failed to upload docs: {response.status_code} - {response.text}"
text = ""
async for chunk in response.aiter_text():
text += chunk
with open(file_path, "w") as f:
f.write(text)
citation_format = r"\[Source:\s*([^,]+?)\s*,\s*Page:\s*(\d+)\s*,\s*Lines:\s*(\d+\s*-\s*\d+)\s*,\s*Start:?\s*(\d+)\]"
assert re.search(citation_format, text) is not None
|