Skip to content

Commit

Permalink
fix custom artifact type in tests
Browse files Browse the repository at this point in the history
  • Loading branch information
connor-mccarthy committed Sep 1, 2022
1 parent 904febf commit 784fed8
Showing 1 changed file with 6 additions and 5 deletions.
11 changes: 6 additions & 5 deletions sdk/python/kfp/components/executor_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,8 @@ def test_input_artifact_custom_type(self):
"""

class VertexDataset:
TYPE_NAME = 'google.VertexDataset'
schema_title = 'google.VertexDataset'
schema_version = '0.0.0'

def __init__(self, name: str, uri: str, metadata: dict) -> None:
self.name = name
Expand All @@ -121,15 +122,15 @@ def __init__(self, name: str, uri: str, metadata: dict) -> None:

@property
def path(self) -> str:
return self.uri.replace('gs://', '/gcs/')
return self.uri.replace('gs://',
artifact_types._GCS_LOCAL_MOUNT_PREFIX)

def test_func(input_artifact_one: Input[VertexDataset]):
self.assertEqual(input_artifact_one.uri,
'gs://some-bucket/input_artifact_one')
self.assertEqual(
input_artifact_one.path,
os.path.join(self._test_dir,
'/gcs/some-bucket/input_artifact_one'))
os.path.join(self._test_dir, 'some-bucket/input_artifact_one'))
self.assertEqual(input_artifact_one.name, 'input_artifact_one')
self.assertIsInstance(input_artifact_one, VertexDataset)

Expand Down Expand Up @@ -832,7 +833,7 @@ def __init__(self, name: str, uri: str, metadata: dict) -> None:

@property
def path(self) -> str:
return self.uri.replace('gs://', '/gcs/')
return self.uri.replace('gs://', artifact_types._GCS_LOCAL_MOUNT_PREFIX)


class TestDictToArtifact(parameterized.TestCase):
Expand Down

0 comments on commit 784fed8

Please sign in to comment.