py: fix missing type check on empty list results

Signed-off-by: Isabella Basso do Amaral <[email protected]>

clients/python/src/model_registry/store/wrapper.py

@@ -177,6 +177,11 @@ def get_contexts(self, ctx_type_name: str, options: ListOptions) -> list[Context
 
         Returns:
             Contexts.
+
+        Raises:
+            TypeNotFoundException: If the type doesn't exist.
+            ServerException: If there was an error getting the type.
+            StoreException: Invalid arguments.
         """
         # TODO: should we make options optional?
         # if options is not None:
@@ -193,9 +198,11 @@ def get_contexts(self, ctx_type_name: str, options: ListOptions) -> list[Context
         # else:
         #     contexts = self._mlmd_store.get_contexts_by_type(ctx_type_name)
 
-        if not contexts:
+        if not contexts and ctx_type_name not in [
+            t.name for t in self._mlmd_store.get_context_types()
+        ]:
             msg = f"Context type {ctx_type_name} does not exist"
-            raise StoreException(msg)
+            raise TypeNotFoundException(msg)
 
         return contexts
 
@@ -316,6 +323,11 @@ def get_artifacts(self, art_type_name: str, options: ListOptions) -> list[Artifa
 
         Returns:
             Artifacts.
+
+        Raises:
+            TypeNotFoundException: If the type doesn't exist.
+            ServerException: If there was an error getting the type.
+            StoreException: Invalid arguments.
         """
         try:
             artifacts = self._mlmd_store.get_artifacts(options)
@@ -327,8 +339,10 @@ def get_artifacts(self, art_type_name: str, options: ListOptions) -> list[Artifa
             raise ServerException(msg) from e
 
         artifacts = self._filter_type(art_type_name, artifacts)
-        if not artifacts:
+        if not artifacts and art_type_name not in [
+            t.name for t in self._mlmd_store.get_artifact_types()
+        ]:
             msg = f"Artifact type {art_type_name} does not exist"
-            raise StoreException(msg)
+            raise TypeNotFoundException(msg)
 
         return artifacts

clients/python/tests/store/test_wrapper.py

@@ -17,6 +17,7 @@
     TypeNotFoundException,
 )
 from model_registry.store import MLMDStore
+from model_registry.types.options import MLMDListOptions
 
 
 @pytest.fixture()
@@ -53,6 +54,28 @@ def test_get_undefined_context_type_id(plain_wrapper: MLMDStore):
         plain_wrapper.get_type_id(Context, "undefined")
 
 
+@pytest.mark.usefixtures("artifact")
+def test_get_no_artifacts(plain_wrapper: MLMDStore):
+    arts = plain_wrapper.get_artifacts("test_artifact", MLMDListOptions())
+    assert arts == []
+
+
+def test_get_undefined_artifacts(plain_wrapper: MLMDStore):
+    with pytest.raises(TypeNotFoundException):
+        plain_wrapper.get_artifacts("undefined", MLMDListOptions())
+
+
+@pytest.mark.usefixtures("context")
+def test_get_no_contexts(plain_wrapper: MLMDStore):
+    ctxs = plain_wrapper.get_contexts("test_context", MLMDListOptions())
+    assert ctxs == []
+
+
+def test_get_undefined_contexts(plain_wrapper: MLMDStore):
+    with pytest.raises(TypeNotFoundException):
+        plain_wrapper.get_contexts("undefined", MLMDListOptions())
+
+
 def test_put_invalid_artifact(plain_wrapper: MLMDStore, artifact: Artifact):
     artifact.properties["null"].int_value = 0