test_pipeline.py

import os
import shutil
import subprocess
from collections.abc import Generator
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import cast
from uuid import uuid4

import psycopg
import pytest
from loguru import logger
from psycopg import Connection, sql
from psycopg.rows import DictRow, TupleRow, dict_row
from testcontainers.core.config import testcontainers_config  # type: ignore

# https://github.com/testcontainers/testcontainers-python/issues/305
from testcontainers.postgres import PostgresContainer  # type: ignore

from constants import IDR_BENE_HISTORY_TABLE
from extractor import PostgresExecutor
from load_events import IdrJobLoadEvent, IdrJobType
from load_partition import LoadType
from load_synthetic import load_from_csv
from logger_config import configure_logger
from model.base_model import LoadMode, Source
from pipeline import run
from pydantic_utils import fields
from settings import enable_prior_auth_ingestion

# ryuk throws a 500 or 404 error for some reason
# seems to have issues with podman https://github.com/testcontainers/testcontainers-python/issues/753
testcontainers_config.ryuk_disabled = True


def _run_migrator(postgres: PostgresContainer) -> None:
    # Python recommends using an absolute path when running an executable
    # to avoid any ambiguity
    mvn = shutil.which("mvn") or "mvn"
    try:
        subprocess.run(
            f"{mvn} flyway:migrate "
            "-Dflyway.url="
            f"jdbc:postgresql://localhost:{postgres.get_exposed_port(5432)}/{postgres.dbname} "
            f"-Dflyway.user={postgres.username} "
            f"-Dflyway.password={postgres.password} "
            "-Duser.timezone=UTC",
            cwd=Path(__file__).parent.joinpath("../bfd-db-migrator-ng"),
            shell=True,
            capture_output=True,
            check=True,
        )
    except subprocess.CalledProcessError as ex:
        print(ex.output)
        raise


def _do_test_pipeline(conn: Connection[DictRow], load_type: LoadType) -> None:
    run(Source.POSTGRES, LoadMode.SYNTHETIC, load_type)

    cur = conn.execute("select * from idr.beneficiary order by bene_sk")
    assert cur.rowcount == 29
    rows = cur.fetchmany(2)

    assert rows[0]["bene_sk"] == 10464258
    assert rows[0]["bene_mbi_id"] == "2ZT2XU2EN18"
    assert rows[1]["bene_sk"] == 16666900
    assert rows[1]["bene_mbi_id"] == "5B88XK5JN88"

    cur = conn.execute("select * from idr.beneficiary_mbi_id order by bene_mbi_id")
    assert cur.rowcount == 24
    rows = cur.fetchmany(1)
    assert rows[0]["bene_mbi_id"] == "1BC3JG0FM51"

    # Xref with valid kill_cred_cd should be included
    cur = conn.execute("select * from idr.beneficiary where bene_sk = 174441863")
    rows = cur.fetchmany(1)
    assert rows[0]["bene_xref_efctv_sk"] == 629529363

    # Xref with no valid entry in v2_bene_xref should not be included
    cur = conn.execute("select * from idr.beneficiary where bene_sk = 353816021")
    rows = cur.fetchmany(1)
    assert rows[0]["bene_xref_efctv_sk"] == 353816021

    if enable_prior_auth_ingestion():
        cur = conn.execute("select * from idr.prior_auth order by mbi_num")
        assert cur.rowcount == 21
        rows = cur.fetchmany(1)
        assert rows[0]["mbi_num"] == "1OX4Y88RV68"

        cur = conn.execute("select * from idr.prior_auth_item order by mbi_num")
        assert cur.rowcount == 64
        rows = cur.fetchmany(1)
        assert rows[0]["mbi_num"] == "1OX4Y88RV68"

    cur = conn.execute("select max(last_ts) as max_ts from idr.load_progress")
    row = cur.fetchone()
    assert row is not None
    max_ts = cast(datetime, row["max_ts"])
    datetime_now = max_ts + timedelta(days=1)
    _advance_time(datetime_now)

    conn.execute(
        f"""
        UPDATE {IDR_BENE_HISTORY_TABLE}
        SET bene_mbi_id = '1S000000000', idr_insrt_ts=%(timestamp)s, idr_updt_ts=%(timestamp)s
        WHERE bene_sk = 10464258
        """,
        {"timestamp": datetime_now},
    )
    conn.commit()

    run(Source.POSTGRES, LoadMode.SYNTHETIC, load_type)

    cur = conn.execute("select * from idr.beneficiary order by bene_sk")
    rows = cur.fetchmany(2)
    assert rows[0]["bene_mbi_id"] == "1S000000000"
    assert rows[1]["bene_mbi_id"] == "5B88XK5JN88"

    cur = conn.execute(
        "select * from idr.beneficiary where bene_kill_cred_cd != '' order by bene_sk"
    )
    assert cur.rowcount == 5
    rows = cur.fetchmany(1)
    assert rows[0]["bene_sk"] == 174441863

    cur = conn.execute("select * from idr.beneficiary_third_party order by bene_sk")
    assert cur.rowcount == 4
    rows = cur.fetchmany(1)
    assert rows[0]["bene_sk"] == 16666900

    cur = conn.execute("select * from idr.beneficiary_status order by bene_sk")
    assert cur.rowcount == 15
    rows = cur.fetchmany(1)
    assert rows[0]["bene_sk"] == 10464258

    cur = conn.execute("select * from idr.beneficiary_entitlement order by bene_sk")
    assert cur.rowcount == 30
    rows = cur.fetchmany(1)
    assert rows[0]["bene_sk"] == 10464258

    cur = conn.execute("select * from idr.beneficiary_entitlement_reason order by bene_sk")
    assert cur.rowcount == 15
    rows = cur.fetchmany(1)
    assert rows[0]["bene_sk"] == 10464258

    cur = conn.execute("select * from idr.beneficiary_dual_eligibility order by bene_sk")
    assert cur.rowcount == 4
    rows = cur.fetchmany(1)
    assert rows[0]["bene_sk"] == 47347082

    cur = conn.execute("select * from idr.beneficiary_overshare_mbi order by bene_mbi_id")
    assert cur.rowcount == 2
    rows = cur.fetchmany(2)
    assert rows[0]["bene_mbi_id"] == "5OH0K85GU23"
    assert rows[1]["bene_mbi_id"] == "6LM1C27GV22"

    cur = conn.execute("select * from idr.contract_pbp_number order by cntrct_pbp_sk")
    assert cur.rowcount == 10
    rows = cur.fetchmany(1)
    assert rows[0]["cntrct_pbp_sk"] == 16513335503

    cur = conn.execute("select * from idr.contract_pbp_contact order by cntrct_pbp_sk")
    assert cur.rowcount == 7
    rows = cur.fetchmany(7)
    assert rows[0]["cntrct_pbp_sk"] == 130640088184
    assert rows[6]["cntrct_pbp_sk"] == 940319838486
    # only a future record exists for this contract
    assert rows[6]["cntrct_pbp_bgn_dt"].strftime("%Y-%m-%d") == "2026-12-01"

    if load_type == LoadType.INITIAL:
        cur = conn.execute("select * from idr.beneficiary_ma_part_d_enrollment order by bene_sk")
        assert cur.rowcount == 4
        rows = cur.fetchmany(1)
        assert rows[0]["bene_sk"] == 353816020
    else:
        cur = conn.execute("select * from idr.beneficiary_ma_part_d_enrollment order by bene_sk")
        assert cur.rowcount == 3
        rows = cur.fetchmany(1)
        assert rows[0]["bene_sk"] == 353816020

    if load_type == LoadType.INITIAL:
        cur = conn.execute("select * from idr.beneficiary_ma_part_d_enrollment_rx order by bene_sk")
        assert cur.rowcount == 3
        rows = cur.fetchmany(1)
        assert rows[0]["bene_sk"] == 353816020
    else:
        cur = conn.execute("select * from idr.beneficiary_ma_part_d_enrollment_rx order by bene_sk")
        assert cur.rowcount == 2
        rows = cur.fetchmany(1)
        assert rows[0]["bene_sk"] == 353816020

    cur = conn.execute("select * from idr.beneficiary_low_income_subsidy order by bene_sk")
    assert cur.rowcount == 2
    rows = cur.fetchmany(1)
    assert rows[0]["bene_sk"] == 353816020

    lis_cmbnd_query = "select * from idr.beneficiary_low_income_subsidy_cmbnd order by bene_sk"
    if load_type == LoadType.INITIAL:
        cur = conn.execute(lis_cmbnd_query)
        assert cur.rowcount == 3
        rows = cur.fetchmany(1)
        assert rows[0]["bene_sk"] == 353816020
    else:
        cur = conn.execute(lis_cmbnd_query)
        assert cur.rowcount == 2
        rows = cur.fetchmany(1)
        assert rows[0]["bene_sk"] == 353816020

    cur = conn.execute("select * from idr.claim_institutional_ss where clm_uniq_id = 8244064276500")
    assert cur.rowcount == 0

    cur = conn.execute("select * from idr.claim_institutional_nch order by clm_uniq_id")
    assert cur.rowcount == 51
    rows = cur.fetchmany(1)
    assert rows[0]["clm_uniq_id"] == 113370100080

    cur = conn.execute("select * from idr.claim_professional_nch order by clm_uniq_id")
    assert cur.rowcount == 51
    rows = cur.fetchmany(1)
    assert rows[0]["clm_uniq_id"] == 119855147698

    cur = conn.execute("select * from idr.claim_professional_ss order by clm_uniq_id")
    assert cur.rowcount == 1
    rows = cur.fetchmany(1)
    assert rows[0]["clm_uniq_id"] == 4991490559710

    cur = conn.execute("select * from idr.claim_rx order by clm_uniq_id")
    assert cur.rowcount == 19
    rows = cur.fetchmany(1)
    assert rows[0]["clm_uniq_id"] == 166776396279

    cur = conn.execute("select * from idr.claim_item_institutional_nch order by clm_uniq_id")
    assert cur.rowcount == 795
    rows = cur.fetchmany(1)
    assert rows[0]["clm_uniq_id"] == 113370100080

    cur = conn.execute("select * from idr.claim_item_professional_nch order by clm_uniq_id")
    assert cur.rowcount == 442
    rows = cur.fetchmany(1)
    assert rows[0]["clm_uniq_id"] == 119855147698

    cur = conn.execute("select * from idr.claim_item_professional_ss order by clm_uniq_id")
    assert cur.rowcount == 1
    rows = cur.fetchmany(1)
    assert rows[0]["clm_uniq_id"] == 4991490559710

    conn.commit()

    # Phase 1 SS (PAC) claims older than 60 days will be pruned on incremental loads
    if load_type == LoadType.INITIAL:
        cur = conn.execute("select * from idr.claim_institutional_ss order by clm_uniq_id")
        assert cur.rowcount == 21
        rows = cur.fetchmany(1)
        assert rows[0]["clm_uniq_id"] == 123359318723

        cur = conn.execute("select * from idr.claim_item_institutional_ss order by clm_uniq_id")
        assert cur.rowcount == 327
        rows = cur.fetchmany(1)
        assert rows[0]["clm_uniq_id"] == 123359318723

    else:
        make_it_stale_ts = datetime.now(UTC) + timedelta(days=60)
        _advance_time(make_it_stale_ts)
        run(Source.POSTGRES, LoadMode.SYNTHETIC, LoadType.INCREMENTAL)
        cur = conn.execute("select * from idr.claim_institutional_ss order by clm_uniq_id")
        assert cur.rowcount == 9
        rows = cur.fetchmany(1)
        assert rows[0]["clm_uniq_id"] == 849348853948

        cur = conn.execute("select * from idr.claim_item_institutional_ss order by clm_uniq_id")
        assert cur.rowcount == 151
        rows = cur.fetchmany(1)
        assert rows[0]["clm_uniq_id"] == 849348853948

    # Test incremental loading logic involving 'source_load_events' if we're testing incremental
    # mode
    if load_type == LoadType.INCREMENTAL:
        # First, pretend that loading ./test_samples1 was the result of loading _all_ possible jobs
        # by inserting load events with completion times of datetime_now + 1hr for all types
        idr_jobs_table = sql.Identifier("idr", "source_load_events")
        cur = conn.execute("select max(last_ts) as max_ts from idr.load_progress")
        row = cur.fetchone()
        assert row is not None
        datetime_now = cast(datetime, row["max_ts"])
        load_1_complete_time = datetime_now + timedelta(hours=1)
        load_jobs = [
            IdrJobLoadEvent(
                id=uuid4(),
                job_type=job_type,
                job_message="SUCCESSFUL",
                event_time=datetime_now,
                completion_time=load_1_complete_time,
            ).model_dump(by_alias=True)
            for job_type in IdrJobType
        ]
        for job in load_jobs:
            conn.execute(
                t"""
                INSERT INTO {idr_jobs_table:i} (
                    {sql.SQL(", ").join(sql.Identifier(k) for k in job):q}
                )
                VALUES (
                    {sql.SQL(", ").join(job.values()):q}
                )
                """
            )
        conn.commit()

        # To simulate a new CLMNCH and FISS load, get a known NCH claim and re-insert it with an
        # updated insert timestamp and ID into the relevant institutional claim staging tables (CLM
        # and CLM_INSTNL)
        staging_clm_table = sql.Identifier("cms_vdm_view_mdcr_prd", "v2_mdcr_clm")
        cur = conn.execute(
            t"""
            SELECT * from {staging_clm_table:i}
            WHERE {"clm_uniq_id":i} = {"0113370100080"}
            """
        )
        conn.commit()
        assert cur.rowcount == 1
        nch_clm_row = cur.fetchmany(1)[0]
        nch_clm_ts = load_1_complete_time + timedelta(hours=1)
        nch_clm_row["clm_uniq_id"] = (
            "9999999999998"  # This clm_uniq_id does not exist in ./test_samples1
        )
        nch_clm_row["clm_num_sk"] = 2
        nch_clm_row["idr_insrt_ts"] = nch_clm_ts
        nch_clm_row["idr_updt_ts"] = nch_clm_ts
        conn.execute(
            t"""
            INSERT INTO {staging_clm_table:i} (
                {sql.SQL(", ").join(sql.Identifier(k) for k in nch_clm_row):q}
            )
            VALUES (
                {sql.SQL(", ").join(nch_clm_row.values()):q}
            )
            """
        )
        conn.commit()
        staging_clm_instnl_table = sql.Identifier("cms_vdm_view_mdcr_prd", "v2_mdcr_clm_instnl")
        cur = conn.execute(
            t"""
            SELECT * from {staging_clm_instnl_table:i}
            WHERE {"clm_dt_sgntr_sk":i} = {"876776550714"}
            """
        )
        conn.commit()
        assert cur.rowcount == 1
        nch_clm_instnl_row = cur.fetchmany(1)[0]
        nch_clm_instnl_row["clm_num_sk"] = 2
        nch_clm_instnl_row["idr_insrt_ts"] = nch_clm_ts
        nch_clm_instnl_row["idr_updt_ts"] = nch_clm_ts
        cur = conn.execute(
            t"""
            INSERT INTO {staging_clm_instnl_table:i} (
                {sql.SQL(", ").join(sql.Identifier(k) for k in nch_clm_instnl_row):q}
            )
            VALUES (
                {sql.SQL(", ").join(nch_clm_instnl_row.values()):q}
            )
            """
        )
        conn.commit()

        # Do it again for a known shared-systems claim
        cur = conn.execute(
            t"""
            SELECT * from {staging_clm_table:i}
            WHERE {"clm_uniq_id":i} = {"849348853948"}
            """
        )
        conn.commit()
        assert cur.rowcount == 1
        ss_clm_row = cur.fetchmany(1)[0]
        ss_clm_ts = load_1_complete_time + timedelta(hours=1)
        ss_clm_row["clm_uniq_id"] = (
            "9999999999999"  # This clm_uniq_id does not exist in ./test_samples1
        )
        ss_clm_row["clm_num_sk"] = 2
        ss_clm_row["idr_insrt_ts"] = ss_clm_ts
        ss_clm_row["idr_updt_ts"] = ss_clm_ts
        conn.execute(
            t"""
            INSERT INTO {staging_clm_table:i} (
                {sql.SQL(", ").join(sql.Identifier(k) for k in ss_clm_row):q}
            )
            VALUES (
                {sql.SQL(", ").join(ss_clm_row.values()):q}
            )
            """
        )
        conn.commit()
        cur = conn.execute(
            t"""
            SELECT * from {staging_clm_instnl_table:i}
            WHERE {"clm_dt_sgntr_sk":i} = {"246326234188"}
            """
        )
        conn.commit()
        assert cur.rowcount == 1
        ss_clm_instnl_row = cur.fetchmany(1)[0]
        ss_clm_instnl_row["clm_num_sk"] = 2
        ss_clm_instnl_row["idr_insrt_ts"] = ss_clm_ts
        ss_clm_instnl_row["idr_updt_ts"] = ss_clm_ts
        cur = conn.execute(
            t"""
            INSERT INTO {staging_clm_instnl_table:i} (
                {sql.SQL(", ").join(sql.Identifier(k) for k in ss_clm_instnl_row):q}
            )
            VALUES (
                {sql.SQL(", ").join(ss_clm_instnl_row.values()):q}
            )
            """
        )
        conn.commit()

        # Simulate running the pipeline in the middle of an "ongoing load" (NCH + SS claims being
        # added)
        _advance_time(ss_clm_ts)
        run(Source.POSTGRES, LoadMode.SYNTHETIC, load_type)

        # Check to make sure the NCH claim was not loaded as no corresponding event should exist
        # in source_load_events nor has it been 24 hours since the last load of NCH data
        nch_table = sql.Identifier("idr", "claim_institutional_nch")
        cur = conn.execute(
            t"""
            SELECT * FROM {nch_table:i}
            WHERE {"clm_uniq_id":i} = {nch_clm_row["clm_uniq_id"]}
            """
        )
        conn.commit()
        assert cur.rowcount == 0

        # _Now_ insert an event into source_load_events indicating that the NCH load job was
        # completed
        nch_load_job = IdrJobLoadEvent(
            id=uuid4(),
            job_type=IdrJobType.NCH,
            job_message="SUCCESSFUL",
            event_time=nch_clm_ts + timedelta(hours=1),
        )
        nch_job_dict = nch_load_job.model_dump(by_alias=True)
        conn.execute(
            t"""
            INSERT INTO {idr_jobs_table:i} (
                {sql.SQL(", ").join(sql.Identifier(k) for k in nch_job_dict):q}
            )
            VALUES (
                {sql.SQL(", ").join(nch_job_dict.values()):q}
            )
            """
        )
        conn.commit()

        # Run the Pipeline with the NCH event having been inserted indicating that there is NCH
        # data to load
        _advance_time(nch_load_job.event_time)
        run(Source.POSTGRES, LoadMode.SYNTHETIC, load_type)

        # Check for the NCH claim in the v3 idr schema
        cur = conn.execute(
            t"""
            SELECT * FROM {nch_table:i}
            WHERE {"clm_uniq_id":i} = {nch_clm_row["clm_uniq_id"]}
            """
        )
        conn.commit()
        assert cur.rowcount == 1
        rows = cur.fetchmany(1)
        assert str(rows[0]["clm_uniq_id"]) == str(nch_clm_row["clm_uniq_id"])

        # Confirm the NCH load event has a completion time
        cur = conn.execute(
            t"""
            SELECT * FROM {idr_jobs_table:i}
            WHERE {fields(IdrJobLoadEvent).id:i} = {nch_load_job.id}
            """
        )
        conn.commit()
        assert cur.rowcount == 1
        updated_nch_job = IdrJobLoadEvent.model_validate(cur.fetchmany(1)[0], by_alias=True)
        assert updated_nch_job.completion_time
        assert updated_nch_job.completion_time >= nch_clm_ts

        # Check that the SS claim was _not_ loaded since its load job has not "yet" completed
        ss_table = sql.Identifier("idr", "claim_institutional_ss")
        cur = conn.execute(
            t"""
            SELECT * FROM {ss_table:i}
            WHERE {"clm_uniq_id":i} = {ss_clm_row["clm_uniq_id"]}
            """
        )
        conn.commit()
        assert cur.rowcount == 0

        # _Now_ insert an event into source_load_events indicating that the FISS load job was
        # completed
        ss_load_job = IdrJobLoadEvent(
            id=uuid4(),
            job_type=IdrJobType.FISS,
            job_message="SUCCESSFUL",
            event_time=ss_clm_ts + timedelta(hours=1.5),
        )
        ss_job_dict = ss_load_job.model_dump(by_alias=True)
        conn.execute(
            t"""
            INSERT INTO {idr_jobs_table:i} (
                {sql.SQL(", ").join(sql.Identifier(k) for k in ss_job_dict):q}
            )
            VALUES (
                {sql.SQL(", ").join(ss_job_dict.values()):q}
            )
            """
        )
        conn.commit()

        # Run one last time now that the FISS "job" has completed and the SS claim can be loaded
        _advance_time(ss_load_job.event_time)
        run(Source.POSTGRES, LoadMode.SYNTHETIC, load_type)

        # Check for the SS claim in the v3 idr schema
        cur = conn.execute(
            t"""
            SELECT * FROM {ss_table:i}
            WHERE {"clm_uniq_id":i} = {ss_clm_row["clm_uniq_id"]}
            """
        )
        conn.commit()
        assert cur.rowcount == 1
        rows = cur.fetchmany(1)
        assert str(rows[0]["clm_uniq_id"]) == str(ss_clm_row["clm_uniq_id"])

        # Confirm the SS load event has a completion time
        cur = conn.execute(
            t"""
            SELECT * FROM {idr_jobs_table:i}
            WHERE {fields(IdrJobLoadEvent).id:i} = {ss_load_job.id}
            """
        )
        conn.commit()
        assert cur.rowcount == 1
        updated_ss_job = IdrJobLoadEvent.model_validate(cur.fetchmany(1)[0], by_alias=True)
        assert updated_ss_job.completion_time
        assert updated_ss_job.completion_time >= ss_clm_ts


def _do_test_prior_auth_update_and_delete(conn: Connection[DictRow], load_type: LoadType) -> None:
    if not enable_prior_auth_ingestion():
        return

    cur = conn.execute(
        "select * from idr.prior_auth where mbi_num = '7ZM6HW2AT68' and utn = '-OTENCJLOQRAKA'"
    )
    assert cur.rowcount == 1
    rows = cur.fetchmany(6)
    assert rows[0]["mbi_num"] == "7ZM6HW2AT68"
    original_updated_ts = rows[0]["bfd_updated_ts"]
    original_name = rows[0]["name"]

    cur = conn.execute(
        "select * from idr.prior_auth where mbi_num = '5OH0K85GU23' and utn = '-SC21YQR4UY4LI'"
    )
    assert cur.rowcount == 1
    row = cur.fetchone()
    assert row is not None

    prauc_table = sql.Identifier("cms_edp_view_cvm_prau_prd", "prauc")
    conn.execute(
        t"""
        UPDATE {prauc_table:i}
        SET name = 'BITE AID PHARMACY'
        WHERE mbi_num = '7ZM6HW2AT68'
        AND utn = '-OTENCJLOQRAKA'
        """
    )

    conn.execute(
        t"""
        DELETE FROM {prauc_table:i}
        WHERE mbi_num = '5OH0K85GU23'
        AND utn = '-SC21YQR4UY4LI'
        """
    )
    conn.commit()

    _advance_time(datetime.now() + timedelta(days=1))
    run(Source.POSTGRES, LoadMode.SYNTHETIC, load_type)

    # verify that updated rows by upstream were updated
    cur = conn.execute(
        "select * from idr.prior_auth where mbi_num = '7ZM6HW2AT68' and utn = '-OTENCJLOQRAKA'"
    )
    assert cur.rowcount == 1
    updated_row = cur.fetchone()
    assert updated_row is not None
    assert updated_row["name"] != original_name
    assert updated_row["bfd_updated_ts"] > original_updated_ts

    # verify that deleted rows by upstream were deleted in header and item level for prior auth
    cur = conn.execute(
        "select * from idr.prior_auth where mbi_num = '5OH0K85GU23' and utn = '-SC21YQR4UY4LI'"
    )
    assert cur.rowcount == 0

    cur = conn.execute(
        "select * from idr.prior_auth_item where mbi_num = '5OH0K85GU23' and utn = '-SC21YQR4UY4LI'"
    )
    assert cur.rowcount == 0

    # verify that untouched rows by upstream were not updated
    cur = conn.execute(
        "select * from idr.prior_auth where mbi_num = '7ZM6HW2AT68' and utn = '-RVUOWAUT5V5QZ'"
    )
    rows = cur.fetchmany(2)
    assert rows[0]["bfd_updated_ts"] < updated_row["bfd_updated_ts"]


def _advance_time(timestamp: datetime) -> None:
    new_time = timestamp + timedelta(minutes=1)
    os.environ["BFD_TEST_DATE"] = new_time.isoformat()


def _reset_db(
    conn: psycopg.Connection[TupleRow], sample_path: Path, postgres: PostgresContainer
) -> None:
    conn.execute(
        """
        DO $$ DECLARE
            r RECORD;
        BEGIN
            FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = 'idr') LOOP
                EXECUTE 'DROP TABLE idr.' || quote_ident(r.tablename) || ' CASCADE';
            END LOOP;

            FOR r IN (
                SELECT tablename FROM pg_tables WHERE schemaname = 'cms_vdm_view_mdcr_prd'
            ) LOOP
                EXECUTE 'DROP TABLE cms_vdm_view_mdcr_prd.'
                    || quote_ident(r.tablename)
                    || ' CASCADE';
            END LOOP;

            FOR r IN (
                SELECT tablename FROM pg_tables 
                WHERE schemaname = 'cms_edp_view_cvm_prau_prd'
            ) LOOP
                EXECUTE 'DROP TABLE cms_edp_view_cvm_prau_prd.'
                    || quote_ident(r.tablename)
                    || ' CASCADE';
            END LOOP;
        END $$;
        """
    )
    conn.commit()

    with Path(__file__).parent.joinpath("./mock-idr.sql").open() as f:
        conn.execute(f.read())  # type: ignore
    conn.commit()

    _run_migrator(postgres)
    load_from_csv(PostgresExecutor(conn), sample_path)  # type: ignore


def _setup_pipeline_environment(info: psycopg.ConnectionInfo) -> None:
    # Info level logs obscure the error output when running tests
    # so we want to override this unless the calling process has set this explicitly
    os.environ.setdefault("IDR_LOG_LEVEL", "warning")
    os.environ["BFD_DB_ENDPOINT"] = info.host
    os.environ["BFD_DB_PORT"] = str(info.port)
    os.environ["BFD_DB_NAME"] = info.dbname
    os.environ["BFD_DB_USERNAME"] = info.user
    os.environ["BFD_DB_PASSWORD"] = info.password
    os.environ["IDR_BATCH_SIZE"] = "100000"
    os.environ["IDR_FORCE_LOAD_PROGRESS"] = "1"
    os.environ["BFD_TEST_DATE"] = "2023-04-02"
    os.environ["IDR_PER_BATCH_MIN_CONNECTIONS"] = "1"
    os.environ["IDR_PER_BATCH_MAX_CONNECTIONS"] = "1"
    os.environ["IDR_ENABLE_PRIOR_AUTH"] = "1"


@pytest.fixture(scope="module")
def postgres_db() -> Generator[tuple[PostgresContainer, str]]:
    with PostgresContainer("postgres:16", driver="") as postgres:
        conninfo = postgres.get_connection_url()
        yield postgres, conninfo


def _test_pipeline_load(postgres_db: tuple[PostgresContainer, str], load_type: LoadType) -> None:
    configure_logger()
    postgres, conninfo = postgres_db
    with psycopg.connect(conninfo=conninfo, row_factory=dict_row) as conn:  # pyright: ignore[reportArgumentType]
        sample_dir = Path(__file__).parent.joinpath("./test_samples1")
        _reset_db(conn, sample_dir, postgres)
        _setup_pipeline_environment(conn.info)
        _do_test_pipeline(cast(Connection[DictRow], conn), load_type)
        _do_test_prior_auth_update_and_delete(cast(Connection[DictRow], conn), load_type)
    logger.remove()


def test_initial_pipeline_load(postgres_db: tuple[PostgresContainer, str]) -> None:
    _test_pipeline_load(postgres_db, LoadType.INITIAL)


def test_incremental_pipeline_load(postgres_db: tuple[PostgresContainer, str]) -> None:
    _test_pipeline_load(postgres_db, LoadType.INCREMENTAL)