feat: add database models (User, GenerationHistory)
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -1 +1,26 @@
|
|||||||
"""History model placeholder."""
|
import uuid
|
||||||
|
from datetime import datetime
|
||||||
|
from sqlalchemy import Column, String, Text, DateTime, ForeignKey, Enum as SQLEnum, Uuid
|
||||||
|
from sqlalchemy.orm import relationship
|
||||||
|
import enum
|
||||||
|
|
||||||
|
from app.database import Base
|
||||||
|
|
||||||
|
class GenerationMode(str, enum.Enum):
|
||||||
|
TEXT2IMAGE = "text2image"
|
||||||
|
IMAGE2IMAGE = "image2image"
|
||||||
|
|
||||||
|
class GenerationHistory(Base):
|
||||||
|
__tablename__ = "generation_history"
|
||||||
|
|
||||||
|
id = Column(Uuid, primary_key=True, default=uuid.uuid4)
|
||||||
|
user_id = Column(Uuid, ForeignKey("users.id"), nullable=False)
|
||||||
|
original_prompt = Column(Text, nullable=False)
|
||||||
|
optimized_prompt = Column(Text, nullable=False)
|
||||||
|
mode = Column(SQLEnum(GenerationMode), nullable=False)
|
||||||
|
model = Column(String(32), nullable=False)
|
||||||
|
image_url = Column(String(512), nullable=False)
|
||||||
|
input_image_url = Column(String(512), nullable=True) # For I2I
|
||||||
|
created_at = Column(DateTime, default=datetime.utcnow)
|
||||||
|
|
||||||
|
user = relationship("User", backref="generation_history")
|
||||||
@@ -1 +1,15 @@
|
|||||||
"""User model placeholder."""
|
import uuid
|
||||||
|
from datetime import datetime
|
||||||
|
from sqlalchemy import Column, String, DateTime, Uuid
|
||||||
|
|
||||||
|
from app.database import Base
|
||||||
|
|
||||||
|
class User(Base):
|
||||||
|
__tablename__ = "users"
|
||||||
|
|
||||||
|
id = Column(Uuid, primary_key=True, default=uuid.uuid4)
|
||||||
|
username = Column(String(32), unique=True, nullable=False, index=True)
|
||||||
|
email = Column(String(255), unique=True, nullable=False, index=True)
|
||||||
|
hashed_password = Column(String(255), nullable=False)
|
||||||
|
created_at = Column(DateTime, default=datetime.utcnow)
|
||||||
|
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
||||||
@@ -0,0 +1,48 @@
|
|||||||
|
import pytest
|
||||||
|
from sqlalchemy import create_engine
|
||||||
|
from sqlalchemy.orm import sessionmaker
|
||||||
|
from app.database import Base
|
||||||
|
from app.models.user import User
|
||||||
|
from app.models.history import GenerationHistory, GenerationMode
|
||||||
|
|
||||||
|
engine = create_engine("sqlite:///:memory:")
|
||||||
|
TestSessionLocal = sessionmaker(bind=engine)
|
||||||
|
|
||||||
|
def test_user_model():
|
||||||
|
Base.metadata.create_all(bind=engine)
|
||||||
|
db = TestSessionLocal()
|
||||||
|
|
||||||
|
user = User(username="testuser", email="test@example.com", hashed_password="hashed")
|
||||||
|
db.add(user)
|
||||||
|
db.commit()
|
||||||
|
|
||||||
|
assert user.id is not None
|
||||||
|
assert user.username == "testuser"
|
||||||
|
assert user.email == "test@example.com"
|
||||||
|
|
||||||
|
db.close()
|
||||||
|
|
||||||
|
def test_generation_history_model():
|
||||||
|
Base.metadata.create_all(bind=engine)
|
||||||
|
db = TestSessionLocal()
|
||||||
|
|
||||||
|
user = User(username="testuser2", email="test2@example.com", hashed_password="hashed")
|
||||||
|
db.add(user)
|
||||||
|
db.commit()
|
||||||
|
|
||||||
|
history = GenerationHistory(
|
||||||
|
user_id=user.id,
|
||||||
|
original_prompt="a cat",
|
||||||
|
optimized_prompt="a fluffy orange cat sitting on a windowsill",
|
||||||
|
mode=GenerationMode.TEXT2IMAGE,
|
||||||
|
model="gemini-2.5-flash-image",
|
||||||
|
image_url="/images/cat.png"
|
||||||
|
)
|
||||||
|
db.add(history)
|
||||||
|
db.commit()
|
||||||
|
|
||||||
|
assert history.id is not None
|
||||||
|
assert history.user_id == user.id
|
||||||
|
assert history.mode == GenerationMode.TEXT2IMAGE
|
||||||
|
|
||||||
|
db.close()
|
||||||
Reference in New Issue
Block a user