From 692ae9ec322fd4fca2fa308d5be6f6dc3a13f7f9 Mon Sep 17 00:00:00 2001
From: Alexey Volkov <alexey.volkov@ark-kun.com>
Date: Wed, 16 Mar 2022 15:07:45 -0700
Subject: [PATCH] tests: Added project ID inference test

---
 .../aiplatform/test_project_id_inference.py   | 91 +++++++++++++++++++
 1 file changed, 91 insertions(+)
 create mode 100644 tests/system/aiplatform/test_project_id_inference.py

diff --git a/tests/system/aiplatform/test_project_id_inference.py b/tests/system/aiplatform/test_project_id_inference.py
new file mode 100644
index 0000000000..292b458d75
--- /dev/null
+++ b/tests/system/aiplatform/test_project_id_inference.py
@@ -0,0 +1,91 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import pytest
+
+from google.cloud import aiplatform
+from google.cloud.aiplatform.compat.types import pipeline_state as gca_pipeline_state
+from tests.system.aiplatform import e2e_base
+
+
+@pytest.mark.usefixtures("prepare_staging_bucket", "delete_staging_bucket")
+class TestProjectIDInference(e2e_base.TestEndToEnd):
+
+    _temp_prefix = "temp-vertex-sdk-project-id-inference"
+
+    def test_project_id_inference(self, shared_state):
+        # Collection of resources generated by this test, to be deleted during teardown
+        shared_state["resources"] = []
+
+        aiplatform.init(
+            location=e2e_base._LOCATION,
+            staging_bucket=shared_state["staging_bucket_name"],
+        )
+
+        worker_pool_specs = [
+            {
+                "machine_spec": {"machine_type": "n1-standard-2",},
+                "replica_count": 1,
+                "container_spec": {
+                    "image_uri": "python:3.9",
+                    "command": [
+                        "sh",
+                        "-exc",
+                        """python3 -m pip install git+https://github.com/Ark-kun/python-aiplatform@fix--Fixed-getitng-project-ID-when-running-on-Vertex-AI#egg=google-cloud-aiplatform&subdirectory=.
+                            "$0" "$@"
+                            """,
+                        "python3",
+                        "-c",
+                        """
+                            from google.cloud import aiplatform
+                            # Not initializing the Vertex SDK explicitly
+                            # Checking teh project ID
+                            print(aiplatform.initializer.global_config.project)
+                            assert not aiplatform.initializer.global_config.project.endswith("-tp")
+                            # Testing ability to list resources
+                            endpoints = aiplatform.Endpoint.list()
+                            print(endpoints)
+                            """,
+                    ],
+                    "args": [],
+                },
+            }
+        ]
+
+        custom_job = aiplatform.CustomJob(
+            display_name=self._make_display_name("custom"),
+            worker_pool_specs=worker_pool_specs,
+        )
+        custom_job.run(
+            enable_web_access=True, sync=False,
+        )
+
+        shared_state["resources"].append(custom_job)
+
+        in_progress_done_check = custom_job.done()
+        custom_job.wait_for_resource_creation()
+
+        completion_done_check = custom_job.done()
+
+        assert (
+            custom_job.state
+            == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED
+        )
+
+        # Check done() method works correctly
+        assert in_progress_done_check is False
+        assert completion_done_check is True