From a935b67e8819767b1de6e3e03fde350f5f8c607a Mon Sep 17 00:00:00 2001
From: zshandy <zshandy@gmail.com>
Date: Mon, 29 May 2023 19:20:33 -0700
Subject: [PATCH] lineage tests

---
 dataprep/{integrate => lineage}/__init__.py   |  0
 dataprep/{integrate => lineage}/lx.py         |  0
 dataprep/tests/lineage/__init__.py            |  0
 .../lineage/dependency_example/a_table.sql    |  2 ++
 .../lineage/dependency_example/aa_table.sql   |  6 ++++
 .../dependency_example/basic_patient_info.sql | 11 +++++++
 .../dependency_example/from_aa_table.sql      |  1 +
 .../lineage/dependency_example/no_dob.sql     |  2 ++
 dataprep/tests/lineage/test_read_sql.py       | 32 +++++++++++++++++++
 9 files changed, 54 insertions(+)
 rename dataprep/{integrate => lineage}/__init__.py (100%)
 rename dataprep/{integrate => lineage}/lx.py (100%)
 create mode 100644 dataprep/tests/lineage/__init__.py
 create mode 100644 dataprep/tests/lineage/dependency_example/a_table.sql
 create mode 100644 dataprep/tests/lineage/dependency_example/aa_table.sql
 create mode 100644 dataprep/tests/lineage/dependency_example/basic_patient_info.sql
 create mode 100644 dataprep/tests/lineage/dependency_example/from_aa_table.sql
 create mode 100644 dataprep/tests/lineage/dependency_example/no_dob.sql
 create mode 100644 dataprep/tests/lineage/test_read_sql.py

diff --git a/dataprep/integrate/__init__.py b/dataprep/lineage/__init__.py
similarity index 100%
rename from dataprep/integrate/__init__.py
rename to dataprep/lineage/__init__.py
diff --git a/dataprep/integrate/lx.py b/dataprep/lineage/lx.py
similarity index 100%
rename from dataprep/integrate/lx.py
rename to dataprep/lineage/lx.py
diff --git a/dataprep/tests/lineage/__init__.py b/dataprep/tests/lineage/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/dataprep/tests/lineage/dependency_example/a_table.sql b/dataprep/tests/lineage/dependency_example/a_table.sql
new file mode 100644
index 000000000..932665b2b
--- /dev/null
+++ b/dataprep/tests/lineage/dependency_example/a_table.sql
@@ -0,0 +1,2 @@
+SELECT subject_id, gender
+FROM `physionet-data.mimiciii_derived.no_dob`;
diff --git a/dataprep/tests/lineage/dependency_example/aa_table.sql b/dataprep/tests/lineage/dependency_example/aa_table.sql
new file mode 100644
index 000000000..ddacce15f
--- /dev/null
+++ b/dataprep/tests/lineage/dependency_example/aa_table.sql
@@ -0,0 +1,6 @@
+CREATE VIEW aa_table AS
+SELECT a.subject_id, b.gender
+FROM `physionet-data.mimiciii_derived.a_table` a, `physionet-data.mimiciii_derived.no_dob` b;
+CREATE TABLE a_table AS
+SELECT subject_id, gender
+FROM `physionet-data.mimiciii_derived.no_dob`;
diff --git a/dataprep/tests/lineage/dependency_example/basic_patient_info.sql b/dataprep/tests/lineage/dependency_example/basic_patient_info.sql
new file mode 100644
index 000000000..4c7f34193
--- /dev/null
+++ b/dataprep/tests/lineage/dependency_example/basic_patient_info.sql
@@ -0,0 +1,11 @@
+-- ------------------------------------------------------------------
+-- Title: Retrieves basic patient information from the patients table
+-- Notes: this query does not specify a schema. To run it on your local
+-- MIMIC schema, run the following command:
+--  SET SEARCH_PATH TO mimiciii;
+-- Where "mimiciii" is the name of your schema, and may be different.
+-- ------------------------------------------------------------------
+
+
+SELECT subject_id, gender, dob
+FROM `physionet-data.mimiciii_clinical.patients`;
diff --git a/dataprep/tests/lineage/dependency_example/from_aa_table.sql b/dataprep/tests/lineage/dependency_example/from_aa_table.sql
new file mode 100644
index 000000000..46bf9ed06
--- /dev/null
+++ b/dataprep/tests/lineage/dependency_example/from_aa_table.sql
@@ -0,0 +1 @@
+SELECT * FROM aa_table;
diff --git a/dataprep/tests/lineage/dependency_example/no_dob.sql b/dataprep/tests/lineage/dependency_example/no_dob.sql
new file mode 100644
index 000000000..79180e2d0
--- /dev/null
+++ b/dataprep/tests/lineage/dependency_example/no_dob.sql
@@ -0,0 +1,2 @@
+SELECT subject_id, gender, dob
+FROM `physionet-data.mimiciii_derived.basic_patient_info`;
diff --git a/dataprep/tests/lineage/test_read_sql.py b/dataprep/tests/lineage/test_read_sql.py
new file mode 100644
index 000000000..31f483e36
--- /dev/null
+++ b/dataprep/tests/lineage/test_read_sql.py
@@ -0,0 +1,32 @@
+# type: ignore
+from os import environ
+import os
+import pytest
+
+from ...lineage import lineagex
+
+
+@pytest.mark.skipif(
+    environ.get("DB_URL", "") == "",
+    reason="Skip tests that requires database setup and sql query specified",
+)
+def test_read_sql() -> None:
+    db_url = environ["DB_URL"]
+    sql = os.path.join(os.getcwd(), 'dependency_example')
+    lx = lineagex(
+        sql,
+        "mimiciii_derived",
+        db_url,
+        "mimiciii_clinical, public"
+    )
+    print("dependency test with database connection", lx)
+    lx = lineagex(
+        sql=sql,
+        target_schema="mimiciii_derived",
+        search_path_schema="mimiciii_clinical, public"
+    )
+    print("dependency test without database connection", lx)
+
+
+if __name__ == "__main__":
+    test_read_sql()