From 9dc6637fff6eeaafc0f98831195c73192338addd Mon Sep 17 00:00:00 2001
From: Eljas Roellin <eljas.roellin@ikmail.com>
Date: Fri, 22 Nov 2024 14:36:37 +0100
Subject: [PATCH] support drug_era

---
 src/ehrdata/io/omop/_check_arguments.py  |  1 +
 src/ehrdata/io/omop/_queries.py          |  3 ++
 tests/data/toy_omop/vanilla/drug_era.csv | 10 ++++
 tests/test_io/test_omop.py               | 62 +++++++++++++++++++++++-
 4 files changed, 75 insertions(+), 1 deletion(-)
 create mode 100644 tests/data/toy_omop/vanilla/drug_era.csv

diff --git a/src/ehrdata/io/omop/_check_arguments.py b/src/ehrdata/io/omop/_check_arguments.py
index 5da1737..b53e79f 100644
--- a/src/ehrdata/io/omop/_check_arguments.py
+++ b/src/ehrdata/io/omop/_check_arguments.py
@@ -17,6 +17,7 @@
     "condition_occurrence",
     "procedure_occurrence",
     "device_exposure",
+    "drug_era",
 ]
 VALID_KEEP_DATES = ["start", "end", "interval"]
 
diff --git a/src/ehrdata/io/omop/_queries.py b/src/ehrdata/io/omop/_queries.py
index 3e15f47..4684633 100644
--- a/src/ehrdata/io/omop/_queries.py
+++ b/src/ehrdata/io/omop/_queries.py
@@ -33,6 +33,7 @@
     "condition_occurrence": "condition",
     "procedure_occurrence": "procedure",
     "device_exposure": "device",
+    "drug_era": "drug",
 }
 # DATA_TABLE_DATE_TRUNK = {
 #     "measurement": "measurement",
@@ -57,6 +58,7 @@
         "condition_occurrence": "condition_start_date",
         "procedure_occurrence": "procedure_date",  # in v5.3, procedure didnt have end date
         "device_exposure": "device_exposure_start_date",
+        "drug_era": "drug_era_start_date",
     },
     "end": {
         "visit_occurrence": "visit_end_date",
@@ -66,6 +68,7 @@
         "condition_occurrence": "condition_end_date",
         "procedure_occurrence": "procedure_end_date",  # in v5.3, procedure didnt have end date TODO v5.3 support
         "device_exposure": "device_exposure_end_date",
+        "drug_era": "drug_era_end_date",
     },
 }
 
diff --git a/tests/data/toy_omop/vanilla/drug_era.csv b/tests/data/toy_omop/vanilla/drug_era.csv
new file mode 100644
index 0000000..a2dbf2c
--- /dev/null
+++ b/tests/data/toy_omop/vanilla/drug_era.csv
@@ -0,0 +1,10 @@
+drug_era_id,person_id,drug_concept_id,drug_era_start_date,drug_era_end_date,drug_exposure_count,gap_days
+1,1,1124957,2100-01-01,2100-01-31,2,1
+2,1,1124957,2100-02-01,2100-02-28,2,1
+3,1,1368671,2100-01-01,2100-01-31,4,3
+4,2,1124957,2100-01-01,2100-01-31,2,1
+5,2,1124957,2100-02-01,2100-02-28,2,1
+6,2,1368671,2100-01-01,2100-01-31,4,3
+7,3,1124957,2100-01-01,2100-01-31,2,2
+8,3,1124957,2100-02-01,2100-02-28,2,1
+9,3,1368671,2100-01-01,2100-01-31,4,3
diff --git a/tests/test_io/test_omop.py b/tests/test_io/test_omop.py
index 911e967..f8ea0af 100644
--- a/tests/test_io/test_omop.py
+++ b/tests/test_io/test_omop.py
@@ -19,6 +19,7 @@
     "condition_occurrence": 2,
     "procedure_occurrence": 2,
     "device_exposure": 2,
+    "drug_era": 2,
 }
 
 # constants for setup_variables
@@ -444,7 +445,66 @@ def test_setup_variables(
                 [[1, 1, 1, 1], [1, 1, 1, 1]],
             ],
         ),
-        # (["note"], ["is_present"]),
+        (
+            ["drug_era"],
+            ["drug_exposure_count"],
+            "start",
+            [
+                [[2, np.nan, np.nan, np.nan], [4, np.nan, np.nan, np.nan]],
+                [[2, np.nan, np.nan, np.nan], [4, np.nan, np.nan, np.nan]],
+                [[2, np.nan, np.nan, np.nan], [4, np.nan, np.nan, np.nan]],
+            ],
+        ),
+        (
+            ["drug_era"],
+            ["drug_exposure_count"],
+            "end",
+            [
+                [[np.nan, np.nan, np.nan, np.nan], [np.nan, np.nan, np.nan, np.nan]],
+                [[np.nan, np.nan, np.nan, np.nan], [np.nan, np.nan, np.nan, np.nan]],
+                [[np.nan, np.nan, np.nan, np.nan], [np.nan, np.nan, np.nan, np.nan]],
+            ],
+        ),
+        (
+            ["drug_era"],
+            ["drug_exposure_count"],
+            "interval",
+            [
+                [[2, 2, 2, 2], [4, 4, 4, 4]],
+                [[2, 2, 2, 2], [4, 4, 4, 4]],
+                [[2, 2, 2, 2], [4, 4, 4, 4]],
+            ],
+        ),
+        (
+            ["drug_era"],
+            ["is_present"],
+            "start",
+            [
+                [[1, np.nan, np.nan, np.nan], [1, np.nan, np.nan, np.nan]],
+                [[1, np.nan, np.nan, np.nan], [1, np.nan, np.nan, np.nan]],
+                [[1, np.nan, np.nan, np.nan], [1, np.nan, np.nan, np.nan]],
+            ],
+        ),
+        (
+            ["drug_era"],
+            ["is_present"],
+            "end",
+            [
+                [[np.nan, np.nan, np.nan, np.nan], [np.nan, np.nan, np.nan, np.nan]],
+                [[np.nan, np.nan, np.nan, np.nan], [np.nan, np.nan, np.nan, np.nan]],
+                [[np.nan, np.nan, np.nan, np.nan], [np.nan, np.nan, np.nan, np.nan]],
+            ],
+        ),
+        (
+            ["drug_era"],
+            ["is_present"],
+            "interval",
+            [
+                [[1, 1, 1, 1], [1, 1, 1, 1]],
+                [[1, 1, 1, 1], [1, 1, 1, 1]],
+                [[1, 1, 1, 1], [1, 1, 1, 1]],
+            ],
+        ),
     ],
 )
 @pytest.mark.parametrize(