309 lines
9.3 KiB
Python
309 lines
9.3 KiB
Python
from __future__ import annotations
|
||
|
||
import base64
|
||
from datetime import UTC, date, datetime
|
||
|
||
import psycopg
|
||
from fastapi.testclient import TestClient
|
||
from PIL import Image
|
||
from psycopg.rows import dict_row
|
||
|
||
from config import Settings
|
||
from web import create_app
|
||
|
||
ADMIN_HEADERS = {
|
||
"Authorization": "Basic "
|
||
+ base64.b64encode(b"admin:secret").decode()
|
||
}
|
||
|
||
|
||
def insert_request(
|
||
database_url: str,
|
||
*,
|
||
user_id: int,
|
||
prompt: str,
|
||
status: str = "done",
|
||
mode: str = "t2i",
|
||
output: str | None = None,
|
||
username: str | None = None,
|
||
enhanced: str | None = None,
|
||
hidden: bool = False,
|
||
) -> int:
|
||
with psycopg.connect(database_url, row_factory=dict_row) as connection:
|
||
row = connection.execute(
|
||
"""
|
||
INSERT INTO requests
|
||
(user_id, username, timestamp, prompt_text, enhanced_prompt,
|
||
status, mode, seed, output_file_path, hidden)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, 1, %s, %s)
|
||
RETURNING id
|
||
""",
|
||
(
|
||
user_id,
|
||
username,
|
||
datetime(2026, 10, 18, 12, tzinfo=UTC),
|
||
prompt,
|
||
enhanced,
|
||
status,
|
||
mode,
|
||
output,
|
||
hidden,
|
||
),
|
||
).fetchone()
|
||
if row is None:
|
||
raise RuntimeError("Не удалось вставить запрос")
|
||
return int(row["id"])
|
||
|
||
|
||
def insert_user(
|
||
database_url: str,
|
||
*,
|
||
user_id: int,
|
||
username: str | None = "alice",
|
||
credits: int = 3,
|
||
is_admin: bool = False,
|
||
is_blocked: bool = False,
|
||
) -> None:
|
||
with psycopg.connect(database_url) as connection:
|
||
connection.execute(
|
||
"""
|
||
INSERT INTO users
|
||
(user_id, username, credits, last_daily_credit, created_at,
|
||
is_admin, is_blocked)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s)
|
||
ON CONFLICT (user_id) DO NOTHING
|
||
""",
|
||
(
|
||
user_id,
|
||
username,
|
||
credits,
|
||
date(2026, 10, 18),
|
||
datetime(2026, 10, 18, 12, tzinfo=UTC),
|
||
is_admin,
|
||
is_blocked,
|
||
),
|
||
)
|
||
|
||
|
||
def test_history_page_and_user_filter(web_settings: Settings) -> None:
|
||
app = create_app(web_settings)
|
||
with TestClient(app) as client:
|
||
insert_request(
|
||
web_settings.database_url,
|
||
user_id=1,
|
||
username="alice",
|
||
prompt="apple prompt",
|
||
)
|
||
insert_request(
|
||
web_settings.database_url,
|
||
user_id=2,
|
||
prompt="banana prompt",
|
||
)
|
||
|
||
response = client.get("/")
|
||
assert response.status_code == 200
|
||
assert "apple prompt" in response.text
|
||
assert "banana prompt" in response.text
|
||
assert "@alice" in response.text
|
||
assert 'title="ID: 2">2</a>' in response.text
|
||
|
||
response = client.get("/?user_id=1")
|
||
assert response.status_code == 200
|
||
assert "apple prompt" in response.text
|
||
assert "banana prompt" not in response.text
|
||
|
||
response = client.get("/?user_id=abc")
|
||
assert response.status_code == 200
|
||
assert "user_id должен состоять только из цифр" in response.text
|
||
|
||
|
||
def test_detail_and_media(web_settings: Settings) -> None:
|
||
output = web_settings.storage_dir / "1" / "1_output.png"
|
||
output.parent.mkdir(parents=True)
|
||
with Image.new("RGB", (320, 200), (180, 40, 40)) as image:
|
||
image.save(output)
|
||
|
||
app = create_app(web_settings)
|
||
with TestClient(app) as client:
|
||
request_id = insert_request(
|
||
web_settings.database_url,
|
||
user_id=1,
|
||
prompt="detail prompt",
|
||
enhanced="An enhanced English prompt",
|
||
output=str(output),
|
||
)
|
||
|
||
detail = client.get(f"/request/{request_id}")
|
||
assert detail.status_code == 200
|
||
assert "detail prompt" in detail.text
|
||
assert "An enhanced English prompt" in detail.text
|
||
|
||
index = client.get("/")
|
||
assert f"/thumb/{request_id}" in index.text
|
||
assert "detail prompt" in index.text
|
||
assert "An enhanced English prompt" in index.text
|
||
|
||
media = client.get(f"/media/{request_id}/output")
|
||
assert media.status_code == 200
|
||
assert media.content.startswith(b"\x89PNG")
|
||
|
||
thumb = client.get(f"/thumb/{request_id}")
|
||
assert thumb.status_code == 200
|
||
assert thumb.headers["content-type"] == "image/jpeg"
|
||
assert thumb.content.startswith(b"\xff\xd8")
|
||
|
||
missing = client.get(f"/media/{request_id}/input")
|
||
assert missing.status_code == 404
|
||
|
||
|
||
def test_user_management(web_settings: Settings) -> None:
|
||
app = create_app(web_settings)
|
||
with TestClient(app) as client:
|
||
insert_user(
|
||
web_settings.database_url, user_id=1, username="alice"
|
||
)
|
||
insert_user(
|
||
web_settings.database_url,
|
||
user_id=2,
|
||
username="bob",
|
||
credits=0,
|
||
)
|
||
|
||
page = client.get("/admin/users", headers=ADMIN_HEADERS)
|
||
assert page.status_code == 200
|
||
assert "@alice" in page.text
|
||
assert "@bob" in page.text
|
||
|
||
assert (
|
||
client.post(
|
||
"/admin/users/1/credits",
|
||
data={"credits": "9"},
|
||
headers=ADMIN_HEADERS,
|
||
follow_redirects=False,
|
||
).status_code
|
||
== 303
|
||
)
|
||
assert (
|
||
client.post(
|
||
"/admin/users/1/admin",
|
||
headers=ADMIN_HEADERS,
|
||
follow_redirects=False,
|
||
).status_code
|
||
== 303
|
||
)
|
||
assert (
|
||
client.post(
|
||
"/admin/users/2/block",
|
||
headers=ADMIN_HEADERS,
|
||
follow_redirects=False,
|
||
).status_code
|
||
== 303
|
||
)
|
||
|
||
with psycopg.connect(
|
||
web_settings.database_url, row_factory=dict_row
|
||
) as connection:
|
||
first = connection.execute(
|
||
"SELECT credits, is_admin, is_blocked FROM users"
|
||
" WHERE user_id = %s",
|
||
(1,),
|
||
).fetchone()
|
||
second = connection.execute(
|
||
"SELECT is_blocked FROM users WHERE user_id = %s", (2,)
|
||
).fetchone()
|
||
assert first is not None
|
||
assert second is not None
|
||
assert first["credits"] == 9
|
||
assert first["is_admin"] is True
|
||
assert first["is_blocked"] is False
|
||
assert second["is_blocked"] is True
|
||
|
||
|
||
def test_admin_auth_and_public_users_redirect(
|
||
web_settings: Settings,
|
||
) -> None:
|
||
app = create_app(web_settings)
|
||
with TestClient(app) as client:
|
||
assert (
|
||
client.get("/admin", follow_redirects=False).status_code == 401
|
||
)
|
||
assert (
|
||
client.get(
|
||
"/admin/users", follow_redirects=False
|
||
).status_code
|
||
== 401
|
||
)
|
||
redirect = client.get("/users", follow_redirects=False)
|
||
assert redirect.status_code == 307
|
||
assert redirect.headers["location"] == "/admin/users"
|
||
assert (
|
||
client.get("/admin", headers=ADMIN_HEADERS).status_code == 200
|
||
)
|
||
assert (
|
||
client.get(
|
||
"/admin/users", headers=ADMIN_HEADERS
|
||
).status_code
|
||
== 200
|
||
)
|
||
|
||
|
||
def test_hidden_requests_are_admin_only(web_settings: Settings) -> None:
|
||
app = create_app(web_settings)
|
||
with TestClient(app) as client:
|
||
visible = insert_request(
|
||
web_settings.database_url,
|
||
user_id=1,
|
||
prompt="visible prompt",
|
||
)
|
||
hidden = insert_request(
|
||
web_settings.database_url,
|
||
user_id=2,
|
||
prompt="hidden prompt",
|
||
hidden=True,
|
||
)
|
||
|
||
public = client.get("/")
|
||
assert "visible prompt" in public.text
|
||
assert "hidden prompt" not in public.text
|
||
assert client.get(f"/request/{hidden}").status_code == 404
|
||
assert client.get(f"/media/{hidden}/output").status_code == 404
|
||
assert client.get(f"/thumb/{hidden}").status_code == 404
|
||
|
||
admin = client.get("/admin", headers=ADMIN_HEADERS)
|
||
assert "hidden prompt" in admin.text
|
||
assert (
|
||
client.get(
|
||
f"/admin/request/{hidden}", headers=ADMIN_HEADERS
|
||
).status_code
|
||
== 200
|
||
)
|
||
|
||
response = client.post(
|
||
f"/admin/requests/{visible}/hide",
|
||
headers=ADMIN_HEADERS,
|
||
follow_redirects=False,
|
||
)
|
||
assert response.status_code == 303
|
||
public = client.get("/")
|
||
assert "visible prompt" not in public.text
|
||
admin = client.get("/admin", headers=ADMIN_HEADERS)
|
||
assert "visible prompt" in admin.text
|
||
|
||
|
||
def test_user_search(web_settings: Settings) -> None:
|
||
app = create_app(web_settings)
|
||
with TestClient(app) as client:
|
||
insert_user(
|
||
web_settings.database_url, user_id=1, username="alice"
|
||
)
|
||
insert_user(
|
||
web_settings.database_url, user_id=2, username="bob"
|
||
)
|
||
|
||
page = client.get(
|
||
"/admin/users?q=alice", headers=ADMIN_HEADERS
|
||
)
|
||
assert page.status_code == 200
|
||
assert "@alice" in page.text
|
||
assert "@bob" not in page.text
|