160 lines
4.6 KiB
Python
160 lines
4.6 KiB
Python
import pytest
|
|
import os
|
|
from fastapi.testclient import TestClient
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
from app.database import Base, get_db
|
|
from app.main import app
|
|
from app.models.user import User
|
|
from app.models.history import GenerationHistory, GenerationMode
|
|
from app.core.security import hash_password, create_access_token
|
|
|
|
TEST_DB_PATH = "/tmp/test_history.db"
|
|
if os.path.exists(TEST_DB_PATH):
|
|
os.remove(TEST_DB_PATH)
|
|
|
|
engine = create_engine(f"sqlite:///{TEST_DB_PATH}", connect_args={"check_same_thread": False})
|
|
TestSessionLocal = sessionmaker(bind=engine)
|
|
|
|
def override_get_db():
|
|
db = TestSessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
app.dependency_overrides[get_db] = override_get_db
|
|
client = TestClient(app)
|
|
|
|
Base.metadata.create_all(bind=engine)
|
|
|
|
def get_auth_header(user_id):
|
|
token = create_access_token({"sub": str(user_id)})
|
|
return {"Authorization": f"Bearer {token}"}
|
|
|
|
def test_get_history_empty():
|
|
db = TestSessionLocal()
|
|
user = User(username="histtest", email="hist@example.com", hashed_password=hash_password("pass"))
|
|
db.add(user)
|
|
db.commit()
|
|
db.refresh(user)
|
|
user_id = user.id
|
|
db.close()
|
|
|
|
response = client.get("/api/history", headers=get_auth_header(user_id))
|
|
assert response.status_code == 200
|
|
assert response.json() == []
|
|
|
|
def test_get_history_with_items():
|
|
db = TestSessionLocal()
|
|
user = User(username="histtest2", email="hist2@example.com", hashed_password=hash_password("pass"))
|
|
db.add(user)
|
|
db.commit()
|
|
db.refresh(user)
|
|
user_id = user.id
|
|
|
|
history = GenerationHistory(
|
|
user_id=user_id,
|
|
original_prompt="a cat",
|
|
optimized_prompt="a fluffy orange cat",
|
|
mode=GenerationMode.TEXT2IMAGE,
|
|
model="gemini",
|
|
image_url="/images/cat.png"
|
|
)
|
|
db.add(history)
|
|
db.commit()
|
|
db.close()
|
|
|
|
response = client.get("/api/history", headers=get_auth_header(user_id))
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert len(data) == 1
|
|
assert data[0]["original_prompt"] == "a cat"
|
|
|
|
def test_get_history_item_by_id():
|
|
db = TestSessionLocal()
|
|
user = User(username="histtest3", email="hist3@example.com", hashed_password=hash_password("pass"))
|
|
db.add(user)
|
|
db.commit()
|
|
db.refresh(user)
|
|
user_id = user.id
|
|
|
|
history = GenerationHistory(
|
|
user_id=user_id,
|
|
original_prompt="a dog",
|
|
optimized_prompt="a fluffy brown dog",
|
|
mode=GenerationMode.TEXT2IMAGE,
|
|
model="gemini",
|
|
image_url="/images/dog.png"
|
|
)
|
|
db.add(history)
|
|
db.commit()
|
|
db.refresh(history)
|
|
history_id = history.id
|
|
db.close()
|
|
|
|
response = client.get(f"/api/history/{history_id}", headers=get_auth_header(user_id))
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["original_prompt"] == "a dog"
|
|
|
|
def test_get_history_item_not_found():
|
|
db = TestSessionLocal()
|
|
user = User(username="histtest4", email="hist4@example.com", hashed_password=hash_password("pass"))
|
|
db.add(user)
|
|
db.commit()
|
|
db.refresh(user)
|
|
user_id = user.id
|
|
db.close()
|
|
|
|
from uuid import uuid4
|
|
fake_id = uuid4()
|
|
|
|
response = client.get(f"/api/history/{fake_id}", headers=get_auth_header(user_id))
|
|
assert response.status_code == 404
|
|
|
|
def test_delete_history_item():
|
|
db = TestSessionLocal()
|
|
user = User(username="histtest5", email="hist5@example.com", hashed_password=hash_password("pass"))
|
|
db.add(user)
|
|
db.commit()
|
|
db.refresh(user)
|
|
user_id = user.id
|
|
|
|
history = GenerationHistory(
|
|
user_id=user_id,
|
|
original_prompt="a bird",
|
|
optimized_prompt="a colorful bird",
|
|
mode=GenerationMode.TEXT2IMAGE,
|
|
model="gemini",
|
|
image_url="/images/bird.png"
|
|
)
|
|
db.add(history)
|
|
db.commit()
|
|
db.refresh(history)
|
|
history_id = history.id
|
|
db.close()
|
|
|
|
response = client.delete(f"/api/history/{history_id}", headers=get_auth_header(user_id))
|
|
assert response.status_code == 204
|
|
|
|
# Verify deleted
|
|
db = TestSessionLocal()
|
|
deleted = db.query(GenerationHistory).filter(GenerationHistory.id == history_id).first()
|
|
db.close()
|
|
assert deleted is None
|
|
|
|
def test_delete_history_item_not_found():
|
|
db = TestSessionLocal()
|
|
user = User(username="histtest6", email="hist6@example.com", hashed_password=hash_password("pass"))
|
|
db.add(user)
|
|
db.commit()
|
|
db.refresh(user)
|
|
user_id = user.id
|
|
db.close()
|
|
|
|
from uuid import uuid4
|
|
fake_id = uuid4()
|
|
|
|
response = client.delete(f"/api/history/{fake_id}", headers=get_auth_header(user_id))
|
|
assert response.status_code == 404 |