Use credentials block in default case

PrefectHQ · Dec 8, 2023 · c84d722 · c84d722
1 parent e58f53c
commit c84d722
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tests/test_aiplatform.py b/tests/test_aiplatform.py
@@ -271,6 +271,7 @@ async def test_generate_work_pool_base_job_template(
     job = VertexAICustomTrainingJob(
         image="docker.io/my_image:latest",
         region="us-central1",
+        gcp_credentials=credentials_block,
     )
     expected_template = default_base_job_template
     default_base_job_template["variables"]["properties"]["image"][
@@ -279,6 +280,9 @@ async def test_generate_work_pool_base_job_template(
     default_base_job_template["variables"]["properties"]["region"][
         "default"
     ] = "us-central1"
+    default_base_job_template["variables"]["properties"]["credentials"]["default"] = {
+        "$ref": {"block_document_id": str(credentials_block._block_document_id)}
+    }
     if job_config == "custom":
         expected_template = base_job_template_with_defaults
         job = VertexAICustomTrainingJob(