Save updated base job template to work pool in provision-infra (PrefectHQ#11355)

serinamarie · web-flow · commit dbf29ad617a5 · 2023-12-07T13:34:33.000-06:00
diff --git a/src/prefect/cli/work_pool.py b/src/prefect/cli/work_pool.py
@@ -434,9 +434,17 @@ async def provision_infrastructure(
                 work_pool.type
             )
             provisioner.console = app.console
-            await provisioner.provision(
+            new_base_job_template = await provisioner.provision(
                 work_pool_name=name, base_job_template=work_pool.base_job_template
             )
+
+            await client.update_work_pool(
+                work_pool_name=name,
+                work_pool=WorkPoolUpdate(
+                    base_job_template=new_base_job_template,
+                ),
+            )
+
         except ValueError as exc:
             app.console.print(f"Error: {exc}")
             app.console.print(
diff --git a/tests/cli/test_work_pool.py b/tests/cli/test_work_pool.py
@@ -797,6 +797,9 @@ async def test_file(self, mock_collection_registry, tmp_path):
 
 class TestProvisionInfrastructure:
     async def test_provision_infra(self, monkeypatch, push_work_pool, prefect_client):
+        client_res = await prefect_client.read_work_pool(push_work_pool.name)
+        assert client_res.base_job_template != FAKE_DEFAULT_BASE_JOB_TEMPLATE
+
         mock_provision = AsyncMock()
 
         class MockProvisioner:
@@ -828,6 +831,10 @@ async def provision(self, *args, **kwargs):
 
         assert mock_provision.await_count == 1
 
+        # ensure work pool base job template was updated
+        client_res = await prefect_client.read_work_pool(push_work_pool.name)
+        assert client_res.base_job_template == FAKE_DEFAULT_BASE_JOB_TEMPLATE
+
     async def test_provision_infra_unsupported(self, push_work_pool):
         res = await run_sync_in_worker_thread(
             invoke_and_assert,
