Add unit tests for dl_database

kbrose · kbrose · commit 42a0dfa6bb7f · 2018-04-17T09:56:22.000-05:00
diff --git a/tests/test_io.py b/tests/test_io.py
@@ -1,6 +1,7 @@
 import wfdb
 import numpy as np
 import os
+import shutil
 
 
 class test_record():
@@ -391,3 +392,20 @@ def tearDownClass(self):
         for file in writefiles:
             if os.path.isfile(file):
                 os.remove(file)
+
+
+class test_download():
+    # Test that we can download records with no "dat" file
+    # Regression test for https://github.com/MIT-LCP/wfdb-python/issues/118
+    def test_dl_database_no_dat_file():
+        wfdb.dl_database('afdb', './download-tests/', ['00735'])
+
+    # Test that we can download records that *do* have a "dat" file.
+    def test_dl_database_with_dat_file():
+        wfdb.dl_database('afdb', './download-tests/', ['04015'])
+
+    # Cleanup written files
+    @classmethod
+    def tearDownClass(self):
+        if os.path.isdir('./download-tests/'):
+            shutil.rmtree('./download-tests/')
