Extra test for artifact listing (#2715)

* extra test for artifacts listing

* return to legacy format

---------

Co-authored-by: Safoine El Khabich <[email protected]>

Author: avishniakov

tests/integration/functional/zen_stores/test_zen_store.py

@@ -2966,6 +2966,40 @@ def test_list_unused_artifacts():
         assert artifact_versions.total == num_unused_artifact_versions_before
 
 
+def test_list_custom_named_artifacts():
+    """Tests listing with `has_custom_name=True` only returns custom named artifacts with proper filtering."""
+    client = Client()
+    store = client.zen_store
+
+    num_artifact_versions_before = store.list_artifact_versions(
+        ArtifactVersionFilter()
+    ).total
+    num_matching_named_before = store.list_artifact_versions(
+        ArtifactVersionFilter(
+            has_custom_name=True, name="contains:test_step_output"
+        )
+    ).total
+    num_runs = 1
+
+    with PipelineRunContext(num_runs):
+        artifact_versions = store.list_artifact_versions(
+            ArtifactVersionFilter()
+        )
+        assert (
+            artifact_versions.total
+            == num_artifact_versions_before + num_runs * 2
+        )
+
+        artifact_versions = store.list_artifact_versions(
+            ArtifactVersionFilter(
+                has_custom_name=True, name="contains:test_step_output"
+            )
+        )
+        assert (
+            artifact_versions.total - num_matching_named_before == num_runs * 2
+        )
+
+
 def test_artifacts_are_not_deleted_with_run(clean_client: "Client"):
     """Tests listing with `unused=True` only returns unused artifacts."""
     store = clean_client.zen_store

tests/integration/functional/zen_stores/utils.py

@@ -17,6 +17,7 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, TypeVar
 
 from pydantic import BaseModel, Field, SecretStr
+from typing_extensions import Annotated
 
 from tests.integration.functional.utils import sample_name
 from zenml import (
@@ -104,13 +105,15 @@
 
 
 @step
-def constant_int_output_test_step() -> int:
+def constant_int_output_test_step() -> Annotated[int, "test_step_output"]:
     logging.info("log")
     return 7
 
 
 @step
-def int_plus_one_test_step(input: int) -> int:
+def int_plus_one_test_step(
+    input: int,
+) -> Annotated[int, "test_step_output_plus_one"]:
     return input + 1