import pytest
from typing import Generator
import psycopg2
from psycopg2.extras import RealDictCursor
from config import settings


@pytest.fixture(scope="session")
def db_connection():
    conn = psycopg2.connect(
        host=settings.DB_HOST,
        port=settings.DB_PORT,
        database=settings.DB_NAME,
        user=settings.DB_USERNAME,
        password=settings.DB_PASSWORD,
        cursor_factory=RealDictCursor
    )
    yield conn
    conn.close()


@pytest.fixture(scope="function")
def db_cursor(db_connection):
    cursor = db_connection.cursor()
    yield cursor
    cursor.close()
    db_connection.rollback()


@pytest.fixture(scope="function")
def clean_test_data(db_cursor):
    yield
    tables = ["sys_user", "sys_role", "sys_menu", "sys_dict", "sys_config"]
    for table in tables:
        db_cursor.execute(f"DELETE FROM {table} WHERE username LIKE 'test_%'")
    db_cursor.connection.commit()