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