from contextlib import contextmanager
from typing import Generator

from aws_testing_utils import lambda_handler, s3_handler
import pytest
from sqlalchemy.orm import scoped_session, Session

from tests.utils import db


@pytest.fixture
def lambda_client() -> lambda_handler:
    return lambda_handler.LambdaHandler()


@pytest.fixture
def s3_client() -> s3_handler:
    return s3_handler.S3Handler()


@contextmanager
def session_scope() -> Generator[Session, None, None]:
    """Provide a scoped transactional session for tests."""
    session_factory: scoped_session[Session] = db.create_session()
    session: Session = session_factory()  # Get the actual Session instance

    try:
        session.expire_all()  # Prevent stale reads
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()
        session_factory.remove()


@pytest.fixture
def db_session() -> Generator[Session, None, None]:
    """Pytest fixture to provide a database session scoped to a test function."""
    with session_scope() as session:
        yield session