From a59f6a60d8a075545217bf8a3a2d73377ad94f78 Mon Sep 17 00:00:00 2001
From: Nargis Sultani <nargis.sultani@cfpb.gov>
Date: Thu, 14 Dec 2023 13:21:49 -0500
Subject: [PATCH] Modified test cases

---
 db_revisions/feed/federal_regulator.csv |  8 +--
 tests/migrations/test_migrations.py     | 66 +++++++++++++++++++------
 2 files changed, 51 insertions(+), 23 deletions(-)

diff --git a/db_revisions/feed/federal_regulator.csv b/db_revisions/feed/federal_regulator.csv
index 43a0273..a7cdab5 100644
--- a/db_revisions/feed/federal_regulator.csv
+++ b/db_revisions/feed/federal_regulator.csv
@@ -5,10 +5,4 @@ FHFA|Federal Housing Finance Agency
 FRS|Federal Reserve System
 NCUA|National Credit Union Administration
 OCC|Office of the Comptroller of the Currency
-OTS|Office of Thrift Supervision (only valid until July 21, 2011)
-
-
-
-
-
-
+OTS|Office of Thrift Supervision (only valid until July 21, 2011)
\ No newline at end of file
diff --git a/tests/migrations/test_migrations.py b/tests/migrations/test_migrations.py
index c021527..8c9bcec 100644
--- a/tests/migrations/test_migrations.py
+++ b/tests/migrations/test_migrations.py
@@ -7,6 +7,9 @@
 import sqlalchemy
 from sqlalchemy import text
 from sqlalchemy.engine import Engine
+import os
+import csv
+from entities.models import AddressStateDao, FederalRegulatorDao, HMDAInstitutionTypeDao, SBLInstitutionTypeDao
 
 from pytest_alembic import MigrationContext
 
@@ -25,36 +28,67 @@ def test_tables_exist_after_migration(alembic_runner: MigrationContext, alembic_
     assert "sbl_institution_type" in tables
 
 
-def test_data_feed_after_migration_045aa502e050(alembic_runner: MigrationContext, alembic_engine: Engine):
+def data_feed_migration_helper(table_name):
+    file_dir = os.path.dirname(os.path.realpath(__file__))
+    data_file_path = f"{file_dir}/../../db_revisions/feed/%s.csv" % table_name
+    with open(data_file_path) as f:
+        reader = csv.reader(f, delimiter="|")
+        next(reader)
+        output_list = list(tuple(line) for line in reader)
+    return output_list
+
+
+def test_address_state_data_feed(alembic_runner: MigrationContext, alembic_engine: Engine):
     # Migrate up to, but not including this new migration
-    alembic_runner.migrate_up_before("045aa502e050")
+    alembic_runner.migrate_up_before("7b6ff51002b5")
 
     # Test address_state feed
-    alembic_runner.insert_into("address_state", dict(code="TS", name="Test State"))
+    address_state_tablename = AddressStateDao.__tablename__
     alembic_runner.migrate_up_one()
     with alembic_engine.connect() as conn:
-        address_state_rows = conn.execute(text("SELECT code from address_state")).fetchall()
-    assert address_state_rows == [("TS",)]
+        address_state_rows = conn.execute(text("SELECT code, name from %s" % address_state_tablename)).fetchall()
+    address_state_expected = data_feed_migration_helper(address_state_tablename)
+    assert address_state_rows == address_state_expected
+
+
+def test_federal_regulator_data_feed(alembic_runner: MigrationContext, alembic_engine: Engine):
+    # Migrate up to, but not including this new migration
+    alembic_runner.migrate_up_before("26a742d97ad9")
 
     # Test federal_regulator feed
-    alembic_runner.insert_into("federal_regulator", dict(id="TFCA", name="Test Farm Credit Administration"))
+    federal_regulator_tablename = FederalRegulatorDao.__tablename__
     alembic_runner.migrate_up_one()
     with alembic_engine.connect() as conn:
-        address_state_rows = conn.execute(text("SELECT id from federal_regulator")).fetchall()
-    assert address_state_rows == [("TFCA",)]
+        federal_regulator_rows = conn.execute(text("SELECT id, name from %s" % federal_regulator_tablename)).fetchall()
+    federal_regulator_expected = data_feed_migration_helper(federal_regulator_tablename)
+    assert federal_regulator_rows == federal_regulator_expected
+
+
+def test_hmda_institution_type_data_feed(alembic_runner: MigrationContext, alembic_engine: Engine):
+    # Migrate up to, but not including this new migration
+    alembic_runner.migrate_up_before("f4ff7d1aa6df")
 
     # Test hmda_institution_type feed
-    alembic_runner.insert_into(
-        "hmda_institution_type", dict(id="T8", name="Test Branch or Agency of FBO (FRS supervised)")
-    )
+    hmda_institution_type_tablename = HMDAInstitutionTypeDao.__tablename__
     alembic_runner.migrate_up_one()
     with alembic_engine.connect() as conn:
-        address_state_rows = conn.execute(text("SELECT id from hmda_institution_type")).fetchall()
-    assert address_state_rows == [("T8",)]
+        hmda_institution_type_rows = conn.execute(
+            text("SELECT id, name from %s" % hmda_institution_type_tablename)
+        ).fetchall()
+    hmda_institution_type_expected = data_feed_migration_helper(hmda_institution_type_tablename)
+    assert hmda_institution_type_rows == hmda_institution_type_expected
+
+
+def test_sbl_institution_type_data_feed(alembic_runner: MigrationContext, alembic_engine: Engine):
+    # Migrate up to, but not including this new migration
+    alembic_runner.migrate_up_before("a41281b1e109")
 
     # Test sbl_institution_type feed
-    alembic_runner.insert_into("sbl_institution_type", dict(id="T12", name="Test Online lender"))
+    sbl_institution_type_tablename = SBLInstitutionTypeDao.__tablename__
     alembic_runner.migrate_up_one()
     with alembic_engine.connect() as conn:
-        address_state_rows = conn.execute(text("SELECT id from sbl_institution_type")).fetchall()
-    assert address_state_rows == [("T12",)]
+        sbl_institution_type_rows = conn.execute(
+            text("SELECT id, name from %s" % sbl_institution_type_tablename)
+        ).fetchall()
+    sbl_institution_type_expected = data_feed_migration_helper(sbl_institution_type_tablename)
+    assert sbl_institution_type_rows == sbl_institution_type_expected