Skip to content

Commit

Permalink
fix: Unbreak aiplatform.Experiment.create (#1509)
Browse files Browse the repository at this point in the history
  • Loading branch information
sasha-gitg authored Jul 15, 2022
1 parent 6ef05de commit 558c141
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 1 deletion.
2 changes: 1 addition & 1 deletion google/cloud/aiplatform/metadata/experiment_resources.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ def create(
credentials=credentials,
)

self = cls.__new__()
self = cls.__new__(cls)
self._metadata_context = experiment_context

return self
Expand Down
26 changes: 26 additions & 0 deletions tests/unit/aiplatform/test_metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -591,6 +591,13 @@ def get_experiment_run_mock():
yield get_context_mock


@pytest.fixture
def create_experiment_context_mock():
with patch.object(MetadataServiceClient, "create_context") as create_context_mock:
create_context_mock.side_effect = [_TEST_EXPERIMENT_CONTEXT]
yield create_context_mock


@pytest.fixture
def create_experiment_run_context_mock():
with patch.object(MetadataServiceClient, "create_context") as create_context_mock:
Expand Down Expand Up @@ -863,6 +870,25 @@ def test_init_experiment_with_existing_metadataStore_and_context(
name=_TEST_CONTEXT_NAME, retry=base._DEFAULT_RETRY
)

@pytest.mark.usefixtures("get_metadata_store_mock")
def test_create_experiment(self, create_experiment_context_mock):
exp = aiplatform.Experiment.create(
experiment_name=_TEST_EXPERIMENT,
description=_TEST_EXPERIMENT_DESCRIPTION,
project=_TEST_PROJECT,
location=_TEST_LOCATION,
)

_TRUE_CONTEXT = copy.deepcopy(_TEST_EXPERIMENT_CONTEXT)
_TRUE_CONTEXT.name = None
_TRUE_CONTEXT.metadata.pop("backing_tensorboard_resource")

create_experiment_context_mock.assert_called_once_with(
parent=_TEST_PARENT, context=_TRUE_CONTEXT, context_id=_TEST_EXPERIMENT
)

assert exp._metadata_context.gca_resource == _TEST_EXPERIMENT_CONTEXT

def test_init_experiment_with_credentials(
self,
get_metadata_store_mock,
Expand Down

0 comments on commit 558c141

Please sign in to comment.