fix: Unbreak aiplatform.Experiment.create (#1509)

sasha-gitg · web-flow · commit 558c1410b433 · 2022-07-15T21:02:17.000Z
diff --git 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
@@ -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,

Original file line number	Diff line number	Diff line change
`@@ -204,7 +204,7 @@ def create(`
`204`	`204`	`credentials=credentials,`
`205`	`205`	`)`
`206`	`206`
`207`		`- self = cls.__new__()`
	`207`	`+ self = cls.__new__(cls)`
`208`	`208`	`self._metadata_context = experiment_context`
`209`	`209`
`210`	`210`	`return self`