567-labs--instructor/examples/sqlmodel/test_basic.py

#!/usr/bin/env python3
"""
Basic SQLModel test to verify core functionality
"""

import logging
from datetime import datetime
from typing import Optional
from uuid import UUID, uuid4
from pydantic import validator
from sqlmodel import Field, SQLModel, create_engine, Session, select, Relationship

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Database setup
engine = create_engine("sqlite:///test_basic.db", echo=False)


# Models with relationships
class Team(SQLModel, table=True):
    """Team model with relationship to heroes"""

    id: Optional[int] = Field(default=None, primary_key=True)
    name: str = Field(min_length=2, max_length=50)
    city: str = Field(min_length=2, max_length=50)
    founded_year: Optional[int] = Field(default=None, ge=1900, le=2024)

    # Relationship to heroes
    heroes: list["Hero"] = Relationship(back_populates="team")


class Hero(SQLModel, table=True):
    """Hero model with auto-generated fields and validation"""

    __table_args__ = {"extend_existing": True}

    # Auto-generated fields
    id: Optional[int] = Field(default=None, primary_key=True)
    created_at: datetime = Field(default_factory=datetime.utcnow)
    uuid: UUID = Field(default_factory=uuid4)

    # Regular fields with validation
    name: str = Field(min_length=2, max_length=50, description="Hero's public name")
    secret_name: str = Field(
        min_length=2, max_length=50, description="Hero's secret identity"
    )
    age: Optional[int] = Field(default=None, ge=16, le=100, description="Hero's age")
    power_level: int = Field(ge=1, le=100, description="Power level from 1-100")
    origin_story: str = Field(
        min_length=10, max_length=200, description="Brief origin story"
    )

    # Foreign key relationship
    team_id: Optional[int] = Field(default=None, foreign_key="team.id")
    team: Optional[Team] = Relationship(back_populates="heroes")

    @validator("name")
    def validate_name_format(cls, v):
        """Ensure hero name doesn't contain inappropriate words"""
        forbidden_words = ["villain", "evil", "bad"]
        if any(word in v.lower() for word in forbidden_words):
            raise ValueError(f"Hero name cannot contain: {', '.join(forbidden_words)}")
        return v


def test_basic_functionality():
    """Test basic SQLModel functionality"""

    # Create tables
    SQLModel.metadata.create_all(engine)
    logger.info("✓ Database tables created")

    # Create a team
    with Session(engine) as session:
        team = Team(name="Avengers", city="New York", founded_year=1963)
        session.add(team)
        session.commit()
        session.refresh(team)
        logger.info(f"✓ Created team: {team.name}")

    # Create heroes
    heroes_data = [
        {
            "name": "Iron Man",
            "secret_name": "Tony Stark",
            "age": 45,
            "power_level": 85,
            "origin_story": "Genius inventor who built a powered suit of armor",
        },
        {
            "name": "Captain America",
            "secret_name": "Steve Rogers",
            "age": 100,
            "power_level": 90,
            "origin_story": "Super soldier enhanced with the super soldier serum",
        },
        {
            "name": "Thor",
            "secret_name": "Thor Odinson",
            "age": 1500,  # This will be clamped to 100 by validation
            "power_level": 95,
            "origin_story": "God of Thunder from Asgard with mystical hammer",
        },
    ]

    created_heroes = []
    with Session(engine) as session:
        # Get the team
        team = session.exec(select(Team).where(Team.name == "Avengers")).first()

        if not team:
            logger.error("Team not found!")
            return

        for hero_data in heroes_data:
            try:
                # Handle age validation
                if hero_data["age"] > 100:
                    hero_data["age"] = 100

                hero = Hero(**hero_data, team_id=team.id)
                session.add(hero)
                created_heroes.append(hero)
                logger.info(f"✓ Created hero: {hero.name}")
            except ValueError as e:
                logger.error(f"✗ Failed to create hero {hero_data['name']}: {e}")

        session.commit()

        # Refresh all heroes
        for hero in created_heroes:
            session.refresh(hero)

    # Test validation
    logger.info("\n--- Testing Validation ---")
    try:
        Hero(
            name="Evil Villain",  # Should trigger validator
            secret_name="Bad Guy",
            power_level=50,
            origin_story="A story of evil deeds",
        )
    except ValueError as e:
        logger.info(f"✓ Validation caught invalid name: {e}")

    # Query with relationships
    logger.info("\n--- Testing Relationships ---")
    with Session(engine) as session:
        # Get team with heroes
        team_with_heroes = session.exec(
            select(Team).where(Team.name == "Avengers")
        ).first()

        if team_with_heroes:
            logger.info(
                f"✓ {team_with_heroes.name} has {len(team_with_heroes.heroes)} heroes"
            )

            for hero in team_with_heroes.heroes:
                logger.info(
                    f"  - {hero.name} ({hero.secret_name}) - Power: {hero.power_level}"
                )

    # Test queries
    logger.info("\n--- Testing Queries ---")
    with Session(engine) as session:
        # Find high-power heroes
        high_power_heroes = session.exec(
            select(Hero).where(Hero.power_level >= 90)
        ).all()

        logger.info(f"✓ Found {len(high_power_heroes)} high-power heroes:")
        for hero in high_power_heroes:
            logger.info(f"  - {hero.name}: {hero.power_level}")

    logger.info("\n✓ All basic functionality tests passed!")


if __name__ == "__main__":
    test_basic_functionality()