From 558c1410b43356da0832b3ded78977ee80be93cd Mon Sep 17 00:00:00 2001
From: sasha-gitg <44654632+sasha-gitg@users.noreply.github.com>
Date: Fri, 15 Jul 2022 17:02:17 -0400
Subject: [PATCH] fix: Unbreak aiplatform.Experiment.create (#1509)

---
 .../metadata/experiment_resources.py          |  2 +-
 tests/unit/aiplatform/test_metadata.py        | 26 +++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

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,