diff --git a/google/cloud/aiplatform/metadata/experiment_resources.py b/google/cloud/aiplatform/metadata/experiment_resources.py index c79de84aa5..e0c48b2f00 100644 --- a/google/cloud/aiplatform/metadata/experiment_resources.py +++ b/google/cloud/aiplatform/metadata/experiment_resources.py @@ -204,7 +204,7 @@ def create( credentials=credentials, ) - self = cls.__new__() + self = cls.__new__(cls) self._metadata_context = experiment_context return self diff --git a/tests/unit/aiplatform/test_metadata.py b/tests/unit/aiplatform/test_metadata.py index fec4642151..514bdb84bd 100644 --- a/tests/unit/aiplatform/test_metadata.py +++ b/tests/unit/aiplatform/test_metadata.py @@ -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: @@ -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,