add custom annotation test

birm · birm · commit f01a151e3f6e · 2022-07-12T14:40:45.000-04:00
diff --git a/tests/test_annotation.py b/tests/test_annotation.py
@@ -3,6 +3,7 @@
 import unittest
 
 import numpy as np
+import pandas as pd
 
 import wfdb
 
@@ -252,13 +253,31 @@ def test_4(self):
         annotation2 = wfdb.rdann("huge", "qrs")
         self.assertEqual(annotation1, annotation2)
 
+    def test_5(self):
+        """
+        Write and read annotations with custom labels.
+        """
+        ann_idx = np.array([1, 1000, 2000, 3000])
+        ann_chan = np.array([3, 1, 2, 3])
+        # write custom labels
+        ann_label_store = np.array([ 4, 2, 1, 3])
+        ann_custom_labels = {'label_store': [1, 2, 3, 4],
+                             'symbol': ['v','l','r','z'],
+                             'description':['pvc','lbbb','rbbb','pac']}
+        ann_custom_labels = pd.DataFrame(data=ann_custom_labels)
+        wfdb.wrann('CustomLabel', 'atr', ann_idx, chan=ann_chan,
+                   custom_labels=ann_custom_labels, label_store=ann_label_store)
+        ann = wfdb.rdann('CustomLabel', 'atr')
+        self.assertEqual(ann.symbol, ['z', 'l', 'v', 'r'])
+
     @classmethod
     def tearDownClass(cls):
         writefiles = [
             "100.atr",
             "1003.atr",
             "12726.anI",
             "huge.qrs",
+            "CustomLabel.atr"
         ]
         for file in writefiles:
             if os.path.isfile(file):
