Skip to content

Commit ffe1cd1

Browse files
author
John Wu
committed
create mimic3 initialize test case with demo dataset
1 parent 6231323 commit ffe1cd1

File tree

1 file changed

+95
-0
lines changed

1 file changed

+95
-0
lines changed

tests/core/test_mimic3.py

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
import unittest
2+
import tempfile
3+
import shutil
4+
import subprocess
5+
import os
6+
from pathlib import Path
7+
8+
from pyhealth.datasets import MIMIC3Dataset
9+
10+
11+
class TestMIMIC3Demo(unittest.TestCase):
12+
"""Test MIMIC3 dataset with demo data downloaded from PhysioNet."""
13+
14+
# Class-level variables for shared dataset
15+
temp_dir = None
16+
demo_dataset_path = None
17+
dataset = None
18+
19+
@classmethod
20+
def setUpClass(cls):
21+
"""Download demo dataset once for all tests."""
22+
cls.temp_dir = tempfile.mkdtemp()
23+
cls._download_demo_dataset()
24+
cls._load_dataset()
25+
26+
@classmethod
27+
def tearDownClass(cls):
28+
"""Clean up downloaded dataset after all tests."""
29+
if cls.temp_dir and os.path.exists(cls.temp_dir):
30+
shutil.rmtree(cls.temp_dir)
31+
32+
@classmethod
33+
def _download_demo_dataset(cls):
34+
"""Download MIMIC-III demo dataset using wget."""
35+
download_url = "https://physionet.org/files/mimiciii-demo/1.4/"
36+
37+
# Use wget to download the demo dataset recursively
38+
cmd = [
39+
"wget",
40+
"-r",
41+
"-N",
42+
"-c",
43+
"-np",
44+
"--directory-prefix",
45+
cls.temp_dir,
46+
download_url,
47+
]
48+
49+
try:
50+
subprocess.run(cmd, check=True, capture_output=True, text=True)
51+
except subprocess.CalledProcessError as e:
52+
raise unittest.SkipTest(f"Failed to download MIMIC-III demo dataset: {e}")
53+
except FileNotFoundError:
54+
raise unittest.SkipTest("wget not available - skipping download test")
55+
56+
# Find the downloaded dataset path
57+
physionet_dir = (
58+
Path(cls.temp_dir) / "physionet.org" / "files" / "mimiciii-demo" / "1.4"
59+
)
60+
if physionet_dir.exists():
61+
cls.demo_dataset_path = str(physionet_dir)
62+
else:
63+
raise unittest.SkipTest("Downloaded dataset not found in expected location")
64+
65+
@classmethod
66+
def _load_dataset(cls):
67+
"""Load the dataset once for all tests."""
68+
tables = ["diagnoses_icd", "procedures_icd", "prescriptions", "noteevents"]
69+
70+
cls.dataset = MIMIC3Dataset(root=cls.demo_dataset_path, tables=tables)
71+
72+
def test_stats(self):
73+
"""Test .stats() method execution."""
74+
try:
75+
self.dataset.stats()
76+
except Exception as e:
77+
self.fail(f"dataset.stats() failed: {e}")
78+
79+
def test_get_events(self):
80+
"""Test get_patient and get_events methods with patient 10006."""
81+
# Test get_patient method
82+
patient = self.dataset.get_patient("10006")
83+
self.assertIsNotNone(patient, msg="Patient 10006 should exist in demo dataset")
84+
85+
# Test get_events method
86+
events = patient.get_events()
87+
self.assertIsNotNone(events, msg="get_events() should not return None")
88+
self.assertIsInstance(events, list, msg="get_events() should return a list")
89+
self.assertGreater(
90+
len(events), 0, msg="get_events() should not return an empty list"
91+
)
92+
93+
94+
if __name__ == "__main__":
95+
unittest.main()

0 commit comments

Comments
 (0)