Skip to content

Commit

Permalink
Extra test for artifact listing (#2715)
Browse files Browse the repository at this point in the history
* extra test for artifacts listing

* return to legacy format

---------

Co-authored-by: Safoine El Khabich <34200873+safoinme@users.noreply.github.com>
  • Loading branch information
avishniakov and safoinme authored May 27, 2024
1 parent 5659da8 commit 574dcd0
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 2 deletions.
34 changes: 34 additions & 0 deletions tests/integration/functional/zen_stores/test_zen_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 5 additions & 2 deletions tests/integration/functional/zen_stores/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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


Expand Down

0 comments on commit 574dcd0

Please sign in to comment.