more test fixes

Emyrk · Emyrk · commit b6c26711fbba · 2025-12-02T11:56:19.000-06:00
diff --git a/enterprise/coderd/provisionerdaemons_test.go b/enterprise/coderd/provisionerdaemons_test.go
@@ -256,21 +256,16 @@ func TestProvisionerDaemonServe(t *testing.T) {
 		authToken := uuid.NewString()
 		data, err := echo.Tar(&echo.Responses{
 			Parse: echo.ParseComplete,
-			ProvisionGraph: []*sdkproto.Response{{
-				Type: &sdkproto.Response_Graph{
-					Graph: &sdkproto.GraphComplete{
-						Resources: []*sdkproto.Resource{{
-							Name: "example",
-							Type: "aws_instance",
-							Agents: []*sdkproto.Agent{{
-								Id:   uuid.NewString(),
-								Name: "example",
-							}},
-						}},
-					},
-				},
-			}},
-			ProvisionGraph: echo.ProvisionGraphWithAgent(authToken),
+			ProvisionGraph: echo.ProvisionGraphWithAgent(authToken, func(g *sdkproto.GraphComplete) {
+				g.Resources = []*sdkproto.Resource{{
+					Name: "example",
+					Type: "aws_instance",
+					Agents: []*sdkproto.Agent{{
+						Id:   uuid.NewString(),
+						Name: "example",
+					}},
+				}}
+			}),
 		})
 		require.NoError(t, err)
 		//nolint:gocritic // Not testing file upload in this test.
diff --git a/provisioner/echo/serve.go b/provisioner/echo/serve.go
@@ -47,22 +47,27 @@ func ProvisionApplyWithAgentAndAPIKeyScope(authToken string, apiKeyScope string)
 
 // ProvisionGraphWithAgent returns provision responses that will mock a fake
 // "aws_instance" resource with an agent that has the given auth token.
-func ProvisionGraphWithAgent(authToken string) []*proto.Response {
+func ProvisionGraphWithAgent(authToken string, muts ...func(g *proto.GraphComplete)) []*proto.Response {
+	gc := &proto.GraphComplete{
+		Resources: []*proto.Resource{{
+			Name: "example",
+			Type: "aws_instance",
+			Agents: []*proto.Agent{{
+				Id:   uuid.NewString(),
+				Name: "example",
+				Auth: &proto.Agent_Token{
+					Token: authToken,
+				},
+			}},
+		}},
+	}
+	for _, mut := range muts {
+		mut(gc)
+	}
+
 	return []*proto.Response{{
 		Type: &proto.Response_Graph{
-			Graph: &proto.GraphComplete{
-				Resources: []*proto.Resource{{
-					Name: "example",
-					Type: "aws_instance",
-					Agents: []*proto.Agent{{
-						Id:   uuid.NewString(),
-						Name: "example",
-						Auth: &proto.Agent_Token{
-							Token: authToken,
-						},
-					}},
-				}},
-			},
+			Graph: gc,
 		},
 	}}
 }
