tests

hdmf-dev · Mar 14, 2024 · 0964adc · 0964adc
1 parent efc30f1
commit 0964adc
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 3 deletions.
diff --git a/src/hdmf/term_set.py b/src/hdmf/term_set.py
@@ -314,8 +314,12 @@ def extend(self, arg):
         This append resolves the wrapper to use the extend of the container using
         the wrapper.
         """
+        if isinstance(arg, np.ndarray):
+            values = arg[self.__field]
+        else:
+            values = [arg]
         bad_data = []
-        for item in arg:
+        for item in values:
             if not self.termset.validate(term=item):
                 bad_data.append(item)
 

diff --git a/tests/unit/common/test_table.py b/tests/unit/common/test_table.py
@@ -231,8 +231,20 @@ def test_add_ref_compound_data_append(self):
         )
         compound_vector_data.append(c_data2)
 
+        np.testing.assert_array_equal(compound_vector_data.data, np.append(c_data, c_data2))
+
     def test_add_ref_compound_data_extend(self):
-        pass
+        c_data = np.array([('Homo sapiens', 24)], dtype=[('species', 'U50'), ('age', 'i4')])
+        c_data2 = np.array([('Mus musculus', 24)], dtype=[('species', 'U50'), ('age', 'i4')])
+        compound_vector_data = VectorData(
+            name='Species_1',
+            description='...',
+            data=c_data
+        )
+        compound_vector_data.extend(c_data2)
+
+        np.testing.assert_array_equal(compound_vector_data.data, np.vstack((c_data, c_data2)))
+
 
     def test_add_ref_wrapped_compound_data_append(self):
         c_data = np.array([('Homo sapiens', 24)], dtype=[('species', 'U50'), ('age', 'i4')])
@@ -245,8 +257,21 @@ def test_add_ref_wrapped_compound_data_append(self):
         )
         compound_vector_data.append(c_data2)
 
+        np.testing.assert_array_equal(compound_vector_data.data.data, np.append(c_data, c_data2))
+
     def test_add_ref_wrapped_compound_data_extend(self):
-        pass
+        c_data = np.array([('Homo sapiens', 24)], dtype=[('species', 'U50'), ('age', 'i4')])
+        c_data2 = np.array([('Mus musculus', 24)], dtype=[('species', 'U50'), ('age', 'i4')])
+        terms = TermSet(term_schema_path='tests/unit/example_test_term_set.yaml')
+        compound_vector_data = VectorData(
+            name='Species_1',
+            description='...',
+            data=TermSetWrapper(value=c_data, field='species', termset=terms)
+        )
+        compound_vector_data.extend(c_data2)
+
+        np.testing.assert_array_equal(compound_vector_data.data.data, np.vstack((c_data, c_data2)))
+
 
     def test_constructor_bad_columns(self):
         columns = ['bad_column']