path: root/test_utils/conftest.py

import os
import shutil
import sys
from datetime import datetime, timedelta, timezone
from os.path import join as pjoin
from pathlib import Path
from typing import TYPE_CHECKING, Callable, Generator
from uuid import uuid4

import django
import pytest
import redis
from _pytest.config import Config
from django.conf import settings as django_settings
from django.core.management import call_command
from dotenv import load_dotenv
from pydantic import MariaDBDsn, PostgresDsn
from redis import Redis

from generalresearch.pg_helper import PostgresConfig
from generalresearch.redis_helper import RedisConfig
from generalresearch.sql_helper import SqlHelper

if TYPE_CHECKING:
    from generalresearch.config import GRLBaseSettings
    from generalresearch.currency import USDCent
    from generalresearch.models.thl.session import Status


@pytest.fixture(scope="session")
def env_file_path(pytestconfig: Config) -> str:
    root_path = pytestconfig.rootpath
    env_file = ".env.test"

    candidates = [
        os.path.join(root_path, env_file),
        os.path.join(root_path, "..", env_file),
    ]

    for env_path in candidates:
        if os.path.exists(env_path):
            load_dotenv(dotenv_path=env_path, override=True)
            return os.path.normpath(env_path)

    raise AssertionError(f"No .env.test file found in: {', '.join(candidates)}")


@pytest.fixture(scope="session")
def settings(env_file_path: str) -> "GRLBaseSettings":
    from generalresearch.config import GRLBaseSettings

    print(f"{env_file_path=}")

    s = GRLBaseSettings(_env_file=env_file_path)

    if s.thl_mkpl_rr_db is not None:
        if s.spectrum_rw_db is None:
            s.spectrum_rw_db = MariaDBDsn(f"{s.thl_mkpl_rw_db}unittest-thl-spectrum")
        if s.spectrum_rr_db is None:
            s.spectrum_rr_db = MariaDBDsn(f"{s.thl_mkpl_rr_db}unittest-thl-spectrum")

    s.mnt_gr_api_dir = pjoin("/tmp", f"test-{uuid4().hex[:12]}")

    return s


# === Database Connectors ===


@pytest.fixture(scope="session")
def postgres_instance(settings: "GRLBaseSettings") -> Generator[PostgresDsn]:
    """Create a ephemeral postgresql instance for us to use during pytest.

    This is simplified, and only based off a single host. We don't want to
    create multiple migrated tmp databases for each rw/rr/ro connection
    """

    assert settings.thl_web_rw_db
    # assert settings.thl_web_rw_db.host

    dsn: PostgresDsn = settings.thl_web_rw_db

    # Connect to default DB to create the new one
    from psycopg import connect
    from psycopg.sql import SQL, Identifier

    now = datetime.now(timezone.utc)
    ts: str = now.strftime("%Y-%m-%d")

    db_name = f"unittest-{ts}-{uuid4().hex[:6]}"
    print("XXX", str(dsn))
    conn = connect(str(dsn))
    conn.autocommit = True
    cur = conn.cursor()
    cur.execute(SQL("CREATE DATABASE {}").format(Identifier(db_name)))
    cur.close()
    conn.close()

    host = dsn.hosts()[0]
    db_url = (
        f"postgres://{host['username']}:{host['password']}@{host['host']}/{db_name}"
    )

    yield PostgresDsn(db_url)

    # Teardown: drop the DB after the session
    conn = connect(str(dsn))
    conn.autocommit = True
    cur = conn.cursor()
    # cur.execute(SQL("DROP DATABASE {}").format(Identifier(db_name)))
    cur.close()
    conn.close()


@pytest.fixture(scope="session")
def django_db_setup(settings: "GRLBaseSettings") -> Callable[..., None]:

    def _inner():

        assert settings.thl_web_rw_db
        dsn: PostgresDsn = settings.thl_web_rw_db
        host = dsn.hosts()[0]

        # 1. Bootstrapping Django settings
        if not django_settings.configured:
            django_settings.configure(
                DATABASES={
                    "default": {
                        "ENGINE": "django.db.backends.postgresql",
                        # PostgresDsn stores path as "/dbname"
                        "NAME": str(dsn.path).lstrip("/"),
                        "USER": host["username"],
                        "PASSWORD": host["password"],
                        "HOST": host["host"],
                        "PORT": "5432",
                    }
                },
                INSTALLED_APPS=[
                    "django.contrib.postgres",
                    "django.contrib.contenttypes",
                    "generalresearch.thl_django",
                ],
            )
        django.setup()

        from django.apps import apps

        for model in apps.get_models():
            print(f"Discovered model: {model._meta.label}")

        # 2. Run migrations directly during fixture activation
        call_command("migrate")

    return _inner


@pytest.fixture(scope="session")
def thl_web_rr(
    settings: "GRLBaseSettings", postgres_instance: PostgresDsn, django_db_setup
) -> PostgresConfig:
    dsn = settings.thl_web_rr_db
    assert dsn
    assert dsn.path

    if dsn.path not in ["/", "/postgres"]:
        assert "/unittest-" in dsn.path

    db_path = postgres_instance.path
    host = dsn.hosts()[0]
    db_url = f"postgres://{host['username']}:{host['password']}@{host['host']}{db_path}"

    # Run Migrations now.
    django_db_setup()

    return PostgresConfig(
        dsn=PostgresDsn(db_url),
        connect_timeout=1,
        statement_timeout=5,
    )


@pytest.fixture(scope="session")
def thl_web_rw(
    settings: "GRLBaseSettings", postgres_instance: PostgresDsn, django_db_setup
) -> PostgresConfig:
    dsn = settings.thl_web_rw_db
    assert dsn
    assert dsn.path

    if dsn.path not in ["/", "/postgres"]:
        assert "/unittest-" in dsn.path

    db_path = postgres_instance.path
    host = dsn.hosts()[0]
    db_url = f"postgres://{host['username']}:{host['password']}@{host['host']}{db_path}"

    # Run Migrations now.
    django_db_setup()

    return PostgresConfig(
        dsn=PostgresDsn(db_url),
        connect_timeout=1,
        statement_timeout=5,
    )


@pytest.fixture(scope="session")
def gr_db(settings: "GRLBaseSettings") -> PostgresConfig:
    dsn = settings.gr_db
    assert dsn
    assert dsn.path

    if dsn.path not in ["/", "/postgres"]:
        assert "/unittest-" in dsn.path

    return PostgresConfig(dsn=settings.gr_db, connect_timeout=5, statement_timeout=2)


@pytest.fixture(scope="session")
def spectrum_rw(settings: "GRLBaseSettings") -> SqlHelper:
    dsn = settings.spectrum_rw_db
    assert dsn
    assert dsn.path

    if dsn.path not in ["/", "/postgres"]:
        assert "/unittest-" in dsn.path

    return SqlHelper(
        dsn=settings.spectrum_rw_db,
        read_timeout=2,
        write_timeout=1,
        connect_timeout=2,
    )


@pytest.fixture(scope="session")
def grliq_db(settings: "GRLBaseSettings") -> PostgresConfig:
    dsn = settings.grliq_db
    assert dsn
    assert dsn.path

    if dsn.path not in ["/", "/postgres"]:
        assert "/unittest-" in dsn.path

    # test_words = {"localhost", "127.0.0.1", "unittest", "grliq-test"}
    # assert any(w in str(postgres_config.dsn) for w in test_words), "check grliq postgres_config"
    # assert "grliqdeceezpocymo" not in str(postgres_config.dsn), "check grliq postgres_config"

    return PostgresConfig(
        dsn=settings.grliq_db,
        connect_timeout=2,
        statement_timeout=2,
    )


@pytest.fixture(scope="session")
def thl_redis(settings: "GRLBaseSettings") -> "Redis":
    # todo: this should get replaced with redisconfig (in most places)
    # I'm not sure where this would be? in the domain name?
    assert "unittest" in str(settings.thl_redis) or "127.0.0.1" in str(
        settings.thl_redis
    )

    return redis.Redis.from_url(
        **{
            "url": str(settings.thl_redis),
            "decode_responses": True,
            "socket_timeout": settings.redis_timeout,
            "socket_connect_timeout": settings.redis_timeout,
        }
    )


@pytest.fixture(scope="session")
def thl_redis_config(settings: "GRLBaseSettings") -> RedisConfig:
    assert "unittest" in str(settings.thl_redis) or "127.0.0.1" in str(
        settings.thl_redis
    )
    return RedisConfig(
        dsn=settings.thl_redis,
        decode_responses=True,
        socket_timeout=settings.redis_timeout,
        socket_connect_timeout=settings.redis_timeout,
    )


@pytest.fixture(scope="session")
def gr_redis_config(settings: "GRLBaseSettings") -> "RedisConfig":
    assert "unittest" in str(settings.gr_redis) or "127.0.0.1" in str(settings.gr_redis)

    return RedisConfig(
        dsn=settings.gr_redis,
        decode_responses=True,
        socket_timeout=settings.redis_timeout,
        socket_connect_timeout=settings.redis_timeout,
    )


@pytest.fixture(scope="session")
def gr_redis(settings: "GRLBaseSettings") -> "Redis":
    assert "unittest" in str(settings.gr_redis) or "127.0.0.1" in str(settings.gr_redis)
    return redis.Redis.from_url(
        **{
            "url": str(settings.gr_redis),
            "decode_responses": True,
            "socket_timeout": settings.redis_timeout,
            "socket_connect_timeout": settings.redis_timeout,
        }
    )


@pytest.fixture
def gr_redis_async(settings: "GRLBaseSettings"):
    assert "unittest" in str(settings.gr_redis) or "127.0.0.1" in str(settings.gr_redis)

    import redis.asyncio as redis_async

    return redis_async.Redis.from_url(
        str(settings.gr_redis),
        decode_responses=True,
        socket_timeout=0.20,
        socket_connect_timeout=0.20,
    )


# === Random helpers ===


@pytest.fixture
def start() -> "datetime":
    from datetime import datetime, timezone

    return datetime(year=1900, month=1, day=1, tzinfo=timezone.utc)


@pytest.fixture
def wall_status(request) -> "Status":
    from generalresearch.models.thl.session import Status

    return request.param if hasattr(request, "wall_status") else Status.COMPLETE


@pytest.fixture
def utc_now() -> datetime:
    return datetime.now(tz=timezone.utc)


@pytest.fixture
def utc_hour_ago() -> datetime:
    return datetime.now(tz=timezone.utc) - timedelta(hours=1)


@pytest.fixture
def utc_day_ago() -> "datetime":
    from datetime import datetime, timedelta, timezone

    return datetime.now(tz=timezone.utc) - timedelta(hours=24)


@pytest.fixture
def utc_90days_ago() -> "datetime":
    from datetime import datetime, timedelta, timezone

    return datetime.now(tz=timezone.utc) - timedelta(days=90)


@pytest.fixture
def utc_60days_ago() -> "datetime":
    from datetime import datetime, timedelta, timezone

    return datetime.now(tz=timezone.utc) - timedelta(days=60)


@pytest.fixture
def utc_30days_ago() -> "datetime":
    from datetime import datetime, timedelta, timezone

    return datetime.now(tz=timezone.utc) - timedelta(days=30)


# === Clean up ===


@pytest.fixture(scope="function")
def delete_df_collection(
    thl_web_rw: PostgresConfig, create_main_accounts: Callable[..., None]
) -> Callable[..., None]:

    from generalresearch.incite.collections import (
        DFCollection,
        DFCollectionType,
    )

    def _inner(coll: "DFCollection"):
        match coll.data_type:
            case DFCollectionType.LEDGER:
                for table in [
                    "ledger_transactionmetadata",
                    "ledger_entry",
                    "ledger_transaction",
                    "ledger_account",
                ]:
                    thl_web_rw.execute_write(
                        query=f"DELETE FROM {table};",
                    )
                create_main_accounts()

            case DFCollectionType.WALL | DFCollectionType.SESSION:
                with thl_web_rw.make_connection() as conn:
                    with conn.cursor() as c:
                        c.execute("SET CONSTRAINTS ALL DEFERRED")
                        for table in [
                            "thl_wall",
                            "thl_session",
                        ]:
                            c.execute(
                                query=f"DELETE FROM {table};",
                            )

            case DFCollectionType.USER:
                for table in ["thl_usermetadata", "thl_user"]:
                    thl_web_rw.execute_write(
                        query=f"DELETE FROM {table};",
                    )

            case _:
                thl_web_rw.execute_write(
                    query=f"DELETE FROM {coll.data_type.value};",
                )

    return _inner


# === GR Related ===


@pytest.fixture(scope="function")
def amount_1(request) -> "USDCent":
    from generalresearch.currency import USDCent

    return USDCent(1)


@pytest.fixture(scope="function")
def amount_100(request) -> "USDCent":
    from generalresearch.currency import USDCent

    return USDCent(100)


def clear_directory(path: Path):
    for entry in os.listdir(path):
        full_path = os.path.join(path, entry)
        if os.path.isfile(full_path) or os.path.islink(full_path):
            os.unlink(full_path)  # remove file or symlink
        elif os.path.isdir(full_path):
            shutil.rmtree(full_path)  # remove folder