Fix assertion in buffer_processor tests

legend-exp · Jan 2, 2024 · c6e3c50 · c6e3c50
1 parent aa0b8f4
commit c6e3c50
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 34 deletions.
diff --git a/tests/buffer_processor/test_buffer_processor.py b/tests/buffer_processor/test_buffer_processor.py
@@ -1039,24 +1039,12 @@ def test_buffer_processor_all_pass(lgnd_test_data, tmptestdir):
     for tb in raw_tables:
         raw, _ = sto.read(tb, raw_file)
         proc, _ = sto.read(tb, processed_file)
-        if isinstance(raw, lgdo.Scalar):
-            raw_value = raw.value
-            raw_attrs = raw.attrs
-            proc_value = proc.value
-            proc_attrs = proc.attrs
-            assert raw_value == proc_value
-            assert raw_attrs == proc_attrs
+
+        if isinstance(raw, lgdo.Struct):
+            for obj in raw:
+                assert raw[obj] == raw[obj]
         else:
-            for obj in raw.keys():
-                if not isinstance(raw[obj], lgdo.Table):
-                    raw_df = raw.view_as("pd", cols=[obj])
-                    proc_df = proc.view_as("pd", cols=[obj])
-                else:
-                    for sub_obj in raw[obj].keys():
-                        raw_df = raw[obj].view_as("pd", cols=[str(sub_obj)])
-                        proc_df = proc[obj].view_as("pd", cols=[str(sub_obj)])
-
-            assert raw_df.equals(proc_df)
+            assert raw == proc
 
 
 # check that packet indexes match in verification test

diff --git a/tests/buffer_processor/test_lh5_buffer_processor.py b/tests/buffer_processor/test_lh5_buffer_processor.py
@@ -1047,24 +1047,12 @@ def test_buffer_processor_all_pass(lgnd_test_data):
     for tb in raw_tables:
         raw, _ = sto.read(tb, raw_file)
         proc, _ = sto.read(tb, processed_file)
-        if isinstance(raw, lgdo.Scalar):
-            raw_value = raw.value
-            raw_attrs = raw.attrs
-            proc_value = proc.value
-            proc_attrs = proc.attrs
-            assert raw_value == proc_value
-            assert raw_attrs == proc_attrs
+
+        if isinstance(raw, lgdo.Struct):
+            for obj in raw:
+                assert raw[obj] == raw[obj]
         else:
-            for obj in raw.keys():
-                if not isinstance(raw[obj], lgdo.Table):
-                    raw_df = raw.view_as("pd", cols=[obj])
-                    proc_df = proc.view_as("pd", cols=[obj])
-                else:
-                    for sub_obj in raw[obj].keys():
-                        raw_df = raw[obj].view_as("pd", cols=[str(sub_obj)])
-                        proc_df = proc[obj].view_as("pd", cols=[str(sub_obj)])
-
-            assert raw_df.equals(proc_df)
+            assert raw == proc
 
 
 def test_lh5_buffer_processor_hdf5_settings(lgnd_test_data):