refactor tests for spike tablizer

AllenInstitute · Nov 20, 2017 · 69ca410 · 69ca410
1 parent dd2b59b
commit 69ca410
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/tests/test_nwb.py b/tests/test_nwb.py
@@ -4,11 +4,21 @@
 
 
 def test_spike_tablizer():
-    compare_list = [ng.nwb.SpikeTablizer().transform({'a':[1.1,2,3], 'b':[1,2.5,6]}),
-                    ng.nwb.SpikeTablizer().fit_transform({'a': [1.1, 2, 3], 'b': [1, 2.5, 6]})]
+    spike_dict = {
+        'a': [1.1,2,3],
+        'b': [1,2.5,6],
+        }
+
+    compare_list = [
+        ng.nwb.SpikeTablizer().transform(spike_dict).reset_index(drop=True),
+        ng.nwb.SpikeTablizer().fit_transform(spike_dict).reset_index(drop=True),
+        ]
+
     base = pd.DataFrame({'neuron': ['b', 'a', 'a', 'b', 'a', 'b'],
-                         'time': [1.0, 1.1, 2.0, 2.5, 3.0, 6.0]}).set_index('time')
+                         'time': [1.0, 1.1, 2.0, 2.5, 3.0, 6.0]})
     for compare in compare_list:
+        print(base)
+        print(compare)
         pdt.assert_frame_equal(base, compare)
 
 if __name__ == "__main__":