From 8fde2ce4441139784bc0fdd62c88d4b833018765 Mon Sep 17 00:00:00 2001
From: sasha-gitg <44654632+sasha-gitg@users.noreply.github.com>
Date: Thu, 26 Aug 2021 15:21:28 -0400
Subject: [PATCH] fix: Populate service_account and network in PipelineJob
 instead of pipeline_spec (#658)

---
 google/cloud/aiplatform/pipeline_jobs.py    | 4 ++--
 tests/unit/aiplatform/test_pipeline_jobs.py | 6 ++++++
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py
index bb98b1f0d5..5c948d17b9 100644
--- a/google/cloud/aiplatform/pipeline_jobs.py
+++ b/google/cloud/aiplatform/pipeline_jobs.py
@@ -250,10 +250,10 @@ def run(
                 Optional. Whether to execute this method synchronously. If False, this method will unblock and it will be executed in a concurrent Future.
         """
         if service_account:
-            self._gca_resource.pipeline_spec.service_account = service_account
+            self._gca_resource.service_account = service_account
 
         if network:
-            self._gca_resource.pipeline_spec.network = network
+            self._gca_resource.network = network
 
         _LOGGER.log_create_with_lro(self.__class__)
 
diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py
index 0e3eddbf22..bc39f4d9d9 100644
--- a/tests/unit/aiplatform/test_pipeline_jobs.py
+++ b/tests/unit/aiplatform/test_pipeline_jobs.py
@@ -84,6 +84,8 @@ def mock_pipeline_service_create():
             name=_TEST_PIPELINE_JOB_NAME,
             state=gca_pipeline_state_v1beta1.PipelineState.PIPELINE_STATE_SUCCEEDED,
             create_time=_TEST_PIPELINE_CREATE_TIME,
+            service_account=_TEST_SERVICE_ACCOUNT,
+            network=_TEST_NETWORK,
         )
         yield mock_create_pipeline_job
 
@@ -93,6 +95,8 @@ def make_pipeline_job(state):
         name=_TEST_PIPELINE_JOB_NAME,
         state=state,
         create_time=_TEST_PIPELINE_CREATE_TIME,
+        service_account=_TEST_SERVICE_ACCOUNT,
+        network=_TEST_NETWORK,
     )
 
 
@@ -239,6 +243,8 @@ def test_run_call_pipeline_service_create(
                 "root": _TEST_PIPELINE_JOB_SPEC["pipelineSpec"]["root"],
             },
             runtime_config=runtime_config,
+            service_account=_TEST_SERVICE_ACCOUNT,
+            network=_TEST_NETWORK,
         )
 
         mock_pipeline_service_create.assert_called_once_with(