From 211718f95a4672fafad1152fead681b90fbba4c6 Mon Sep 17 00:00:00 2001
From: Spike Curtis <spike@coder.com>
Date: Tue, 10 Oct 2023 13:24:43 +0400
Subject: [PATCH] fix: fix MaliciousTar test case (#10158)

fixes #9895

Problem was that provisionerd tries to acquire the next job, and races with shutdown, triggering the assert in the handler.  Switches this test case to use the more robust handler.
---
 provisionerd/provisionerd_test.go | 31 ++++++++++++++-----------------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/provisionerd/provisionerd_test.go b/provisionerd/provisionerd_test.go
index 3d8617a4a5..c39edd491a 100644
--- a/provisionerd/provisionerd_test.go
+++ b/provisionerd/provisionerd_test.go
@@ -146,27 +146,24 @@ func TestProvisionerd(t *testing.T) {
 		var (
 			completeChan = make(chan struct{})
 			completeOnce sync.Once
+			acq          = newAcquireOne(t, &proto.AcquiredJob{
+				JobId:       "test",
+				Provisioner: "someprovisioner",
+				TemplateSourceArchive: createTar(t, map[string]string{
+					"../../../etc/passwd": "content",
+				}),
+				Type: &proto.AcquiredJob_TemplateImport_{
+					TemplateImport: &proto.AcquiredJob_TemplateImport{
+						Metadata: &sdkproto.Metadata{},
+					},
+				},
+			})
 		)
 
 		closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
 			return createProvisionerDaemonClient(t, done, provisionerDaemonTestServer{
-				acquireJobWithCancel: func(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error {
-					err := stream.Send(&proto.AcquiredJob{
-						JobId:       "test",
-						Provisioner: "someprovisioner",
-						TemplateSourceArchive: createTar(t, map[string]string{
-							"../../../etc/passwd": "content",
-						}),
-						Type: &proto.AcquiredJob_TemplateImport_{
-							TemplateImport: &proto.AcquiredJob_TemplateImport{
-								Metadata: &sdkproto.Metadata{},
-							},
-						},
-					})
-					assert.NoError(t, err)
-					return nil
-				},
-				updateJob: noopUpdateJob,
+				acquireJobWithCancel: acq.acquireWithCancel,
+				updateJob:            noopUpdateJob,
 				failJob: func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error) {
 					completeOnce.Do(func() { close(completeChan) })
 					return &proto.Empty{}, nil