#   Copyright 2021 Modelyst LLC
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import logging
import os
import sys

import pytest
from psycopg import connect as pg3_connect
from sqlalchemy import MetaData
from sqlmodel import Session, create_engine, text

from dbgen.configuration import config
from dbgen.core.entity import BaseEntity
from dbgen.core.metadata import meta_registry


@pytest.fixture()
def clear_registry():
    # Clear the tables in the metadata for the default base model
    BaseEntity.metadata.clear()
    # Clear the Models associated with the registry, to avoid warnings
    BaseEntity._sa_registry.dispose()
    yield
    BaseEntity.metadata.clear()
    BaseEntity._sa_registry.dispose()


@pytest.fixture(scope="module")
def sql_engine():
    dsn = os.environ.get('TEST_DSN', config.main_dsn)
    engine = create_engine(dsn)
    return engine


@pytest.fixture(scope="function")
def connection(sql_engine):
    """sql_engine connection"""
    metadata = MetaData()
    metadata.reflect(sql_engine)
    metadata.drop_all(sql_engine)
    connection = sql_engine.connect()
    yield connection
    connection.close()


@pytest.fixture(scope="function")
def session(connection):
    transaction = connection.begin()
    session = Session(bind=connection, autocommit=False, autoflush=True)
    yield session
    transaction.rollback()
    transaction.close()
    session.close()


@pytest.fixture(scope="function")
def seed_db(connection):
    connection.execute(text("CREATE table users (id serial primary key, name text);"))
    for user in range(100):
        connection.execute(text(f"INSERT into users(name) values ('user_{user}');"))
    connection.commit()
    yield
    connection.execute(text("drop table users;"))
    connection.commit()


@pytest.fixture(scope="function")
def make_db(connection):
    pass

    metadata = MetaData()
    metadata.reflect(connection)
    metadata.drop_all(connection)
    BaseEntity.metadata.create_all(connection)
    connection.commit()
    yield
    BaseEntity.metadata.drop_all(connection)
    connection.commit()


@pytest.fixture(scope="function")
def raw_connection(make_db, sql_engine):
    raw = sql_engine.raw_connection()
    yield raw
    raw.close()


@pytest.fixture(scope="function")
def raw_pg3_connection(make_db, sql_engine):
    connection = pg3_connect(str(sql_engine.url))
    yield connection
    connection.close()


@pytest.fixture
def debug_logger():
    custom_logger = logging.getLogger("dbgen")
    custom_logger.propagate = True
    custom_logger.setLevel(logging.DEBUG)

    log_format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s Test"
    formatter = logging.Formatter(log_format)
    console_handler = logging.StreamHandler(stream=sys.stdout)
    console_handler.setFormatter(formatter)
    custom_logger.addHandler(console_handler)
    return custom_logger


@pytest.fixture(scope='function')
def recreate_meta(connection):
    connection.execute(text(f'create schema if not exists {config.meta_schema}'))
    meta_registry.metadata.drop_all(connection)
    meta_registry.metadata.create_all(connection)
    yield

sqlalchemy.MetaData dbgen.core.metadata.meta_registry.metadata.drop_all dbgen.core.metadata.meta_registry.metadata.create_all sqlmodel.create_engine sqlmodel.Session logging.StreamHandler pytest.fixture sqlmodel.text logging.getLogger dbgen.core.entity.BaseEntity.metadata.create_all os.environ.get logging.Formatter dbgen.core.entity.BaseEntity.metadata.drop_all dbgen.core.entity.BaseEntity.metadata.clear dbgen.core.entity.BaseEntity._sa_registry.dispose
[((903, 919), 'pytest.fixture', 'pytest.fixture', ([], {}), '()\n', (917, 919), False, 'import pytest\n'), ((1232, 1262), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""module"""'}), "(scope='module')\n", (1246, 1262), False, 'import pytest\n'), ((1388, 1420), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""function"""'}), "(scope='function')\n", (1402, 1420), False, 'import pytest\n'), ((1659, 1691), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""function"""'}), "(scope='function')\n", (1673, 1691), False, 'import pytest\n'), ((1919, 1951), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""function"""'}), "(scope='function')\n", (1933, 1951), False, 'import pytest\n'), ((2288, 2320), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""function"""'}), "(scope='function')\n", (2302, 2320), False, 'import pytest\n'), ((2602, 2634), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""function"""'}), "(scope='function')\n", (2616, 2634), False, 'import pytest\n'), ((2747, 2779), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""function"""'}), "(scope='function')\n", (2761, 2779), False, 'import pytest\n'), ((3386, 3418), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""function"""'}), "(scope='function')\n", (3400, 3418), False, 'import pytest\n'), ((1012, 1039), 'dbgen.core.entity.BaseEntity.metadata.clear', 'BaseEntity.metadata.clear', ([], {}), '()\n', (1037, 1039), False, 'from dbgen.core.entity import BaseEntity\n'), ((1115, 1148), 'dbgen.core.entity.BaseEntity._sa_registry.dispose', 'BaseEntity._sa_registry.dispose', ([], {}), '()\n', (1146, 1148), False, 'from dbgen.core.entity import BaseEntity\n'), ((1163, 1190), 'dbgen.core.entity.BaseEntity.metadata.clear', 'BaseEntity.metadata.clear', ([], {}), '()\n', (1188, 1190), False, 'from dbgen.core.entity import BaseEntity\n'), ((1195, 1228), 'dbgen.core.entity.BaseEntity._sa_registry.dispose', 'BaseEntity._sa_registry.dispose', ([], {}), '()\n', (1226, 1228), False, 'from dbgen.core.entity import BaseEntity\n'), ((1291, 1334), 'os.environ.get', 'os.environ.get', (['"""TEST_DSN"""', 'config.main_dsn'], {}), "('TEST_DSN', config.main_dsn)\n", (1305, 1334), False, 'import os\n'), ((1348, 1366), 'sqlmodel.create_engine', 'create_engine', (['dsn'], {}), '(dsn)\n', (1361, 1366), False, 'from sqlmodel import Session, create_engine, text\n'), ((1496, 1506), 'sqlalchemy.MetaData', 'MetaData', ([], {}), '()\n', (1504, 1506), False, 'from sqlalchemy import MetaData\n'), ((1768, 1826), 'sqlmodel.Session', 'Session', ([], {'bind': 'connection', 'autocommit': '(False)', 'autoflush': '(True)'}), '(bind=connection, autocommit=False, autoflush=True)\n', (1775, 1826), False, 'from sqlmodel import Session, create_engine, text\n'), ((2371, 2381), 'sqlalchemy.MetaData', 'MetaData', ([], {}), '()\n', (2379, 2381), False, 'from sqlalchemy import MetaData\n'), ((2453, 2495), 'dbgen.core.entity.BaseEntity.metadata.create_all', 'BaseEntity.metadata.create_all', (['connection'], {}), '(connection)\n', (2483, 2495), False, 'from dbgen.core.entity import BaseEntity\n'), ((2534, 2574), 'dbgen.core.entity.BaseEntity.metadata.drop_all', 'BaseEntity.metadata.drop_all', (['connection'], {}), '(connection)\n', (2562, 2574), False, 'from dbgen.core.entity import BaseEntity\n'), ((2977, 3003), 'logging.getLogger', 'logging.getLogger', (['"""dbgen"""'], {}), "('dbgen')\n", (2994, 3003), False, 'import logging\n'), ((3175, 3204), 'logging.Formatter', 'logging.Formatter', (['log_format'], {}), '(log_format)\n', (3192, 3204), False, 'import logging\n'), ((3227, 3267), 'logging.StreamHandler', 'logging.StreamHandler', ([], {'stream': 'sys.stdout'}), '(stream=sys.stdout)\n', (3248, 3267), False, 'import logging\n'), ((3536, 3579), 'dbgen.core.metadata.meta_registry.metadata.drop_all', 'meta_registry.metadata.drop_all', (['connection'], {}), '(connection)\n', (3567, 3579), False, 'from dbgen.core.metadata import meta_registry\n'), ((3584, 3629), 'dbgen.core.metadata.meta_registry.metadata.create_all', 'meta_registry.metadata.create_all', (['connection'], {}), '(connection)\n', (3617, 3629), False, 'from dbgen.core.metadata import meta_registry\n'), ((2000, 2062), 'sqlmodel.text', 'text', (['"""CREATE table users (id serial primary key, name text);"""'], {}), "('CREATE table users (id serial primary key, name text);')\n", (2004, 2062), False, 'from sqlmodel import Session, create_engine, text\n'), ((2234, 2259), 'sqlmodel.text', 'text', (['"""drop table users;"""'], {}), "('drop table users;')\n", (2238, 2259), False, 'from sqlmodel import Session, create_engine, text\n'), ((3473, 3530), 'sqlmodel.text', 'text', (['f"""create schema if not exists {config.meta_schema}"""'], {}), "(f'create schema if not exists {config.meta_schema}')\n", (3477, 3530), False, 'from sqlmodel import Session, create_engine, text\n'), ((2119, 2175), 'sqlmodel.text', 'text', (['f"""INSERT into users(name) values (\'user_{user}\');"""'], {}), '(f"INSERT into users(name) values (\'user_{user}\');")\n', (2123, 2175), False, 'from sqlmodel import Session, create_engine, text\n')]