fix testcases

georgia-tech-db · Oct 17, 2023 · c5a66f6 · c5a66f6
1 parent d9a8359
commit c5a66f6
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 9 deletions.
diff --git a/evadb/catalog/services/function_catalog_service.py b/evadb/catalog/services/function_catalog_service.py
@@ -13,10 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from typing import List
+
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.expression import select
 
 from evadb.catalog.models.function_catalog import FunctionCatalog, FunctionCatalogEntry
+from evadb.catalog.models.utils import (
+    FunctionIOCatalogEntry,
+    FunctionMetadataCatalogEntry,
+)
 from evadb.catalog.services.base_service import BaseService
 from evadb.catalog.services.function_io_catalog_service import FunctionIOCatalogService
 from evadb.catalog.services.function_metadata_catalog_service import (
@@ -38,8 +43,8 @@ def insert_entry(
         impl_path: str,
         type: str,
         checksum: str,
-        function_io_list: List["FunctionIOCatalogEntry"],
-        function_metadata_list: List["FunctionMetadataCatalogEntry"],
+        function_io_list: List[FunctionIOCatalogEntry],
+        function_metadata_list: List[FunctionMetadataCatalogEntry],
     ) -> FunctionCatalogEntry:
         """Insert a new function entry
 
@@ -56,10 +61,10 @@ def insert_entry(
         function_obj = function_obj.save(self.session)
 
         for function_io in function_io_list:
-            function_io.function_id = function_obj.row_id
+            function_io.function_id = function_obj._row_id
         io_objs = self._function_io_service.create_entries(function_io_list)
         for function_metadata in function_metadata_list:
-            function_metadata.function_id = function_obj.row_id
+            function_metadata.function_id = function_obj._row_id
         metadata_objs = self._function_metadata_service.create_entries(
             function_metadata_list
         )

diff --git a/evadb/catalog/services/function_io_catalog_service.py b/evadb/catalog/services/function_io_catalog_service.py
@@ -82,3 +82,4 @@ def create_entries(self, io_list: List[FunctionIOCatalogEntry]):
                 function_id=io.function_id,
             )
             io_objs.append(io_obj)
+        return io_objs
diff --git a/evadb/catalog/services/function_metadata_catalog_service.py b/evadb/catalog/services/function_metadata_catalog_service.py
@@ -38,6 +38,7 @@ def create_entries(self, entries: List[FunctionMetadataCatalogEntry]):
                     key=entry.key, value=entry.value, function_id=entry.function_id
                 )
                 metadata_objs.append(metadata_obj)
+            return metadata_objs
         except Exception as e:
             raise CatalogError(e)
 

diff --git a/test/unit_tests/catalog/test_catalog_manager.py b/test/unit_tests/catalog/test_catalog_manager.py
@@ -142,12 +142,13 @@ def test_insert_function(
             function_io_list,
             function_metadata_list,
         )
-        functionio_mock.return_value.insert_entries.assert_called_with(function_io_list)
-        functionmetadata_mock.return_value.insert_entries.assert_called_with(
-            function_metadata_list
-        )
         function_mock.return_value.insert_entry.assert_called_with(
-            "function", "sample.py", "classification", checksum_mock.return_value
+            "function",
+            "sample.py",
+            "classification",
+            checksum_mock.return_value,
+            function_io_list,
+            function_metadata_list,
         )
         checksum_mock.assert_called_with("sample.py")
         self.assertEqual(actual, function_mock.return_value.insert_entry.return_value)