From 0964adcc2021c80c5f412660fe3b7597a870402d Mon Sep 17 00:00:00 2001
From: mavaylon1 <mavaylon1@berkeley.edu>
Date: Thu, 14 Mar 2024 13:40:06 -0700
Subject: [PATCH] tests

---
 src/hdmf/term_set.py            |  6 +++++-
 tests/unit/common/test_table.py | 29 +++++++++++++++++++++++++++--
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/src/hdmf/term_set.py b/src/hdmf/term_set.py
index 8308393d5..abbaf0b3c 100644
--- 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
index 11cce7ed8..e9e3d18c3 100644
--- 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']