tests/utils.py - master - tsps/core

from sqlalchemy import text
from typing import Type, Any

from app.database import SessionLocal


def clean_table(model: Type[Any]):
    """
    Очищает таблицу указанной ORM-модели.
    Работает только в PostgreSQL. Сбрасывает счётчики автоинкремента.
    """
    db = SessionLocal()
    try:
        table_name = model.__tablename__
        # Используем TRUNCATE с сбросом автоинкремента
        db.execute(text(f'TRUNCATE TABLE "{table_name}" RESTART IDENTITY CASCADE'))
        db.commit()
    except Exception as e:
        db.rollback()
        raise e
    finally:
        db.close()