chore: fix TestManagedAgentLimit flake

deansheather · deansheather · commit 470b402f1144 · 2025-07-24T04:54:00.000Z
diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -830,7 +830,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 						}
 						api.derpMesh.SetAddresses(addresses, false)
 					}
-					_ = api.updateEntitlements(ctx)
+					_ = api.updateEntitlements(api.ctx)
 				})
 			} else {
 				coordinator = agpltailnet.NewCoordinator(api.Logger)
@@ -840,7 +840,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 				api.replicaManager.SetCallback(func() {
 					// If the amount of replicas change, so should our entitlements.
 					// This is to display a warning in the UI if the user is unlicensed.
-					_ = api.updateEntitlements(ctx)
+					_ = api.updateEntitlements(api.ctx)
 				})
 			}
 
diff --git a/enterprise/coderd/license/license.go b/enterprise/coderd/license/license.go
@@ -432,7 +432,11 @@ func LicensesEntitlements(
 		if featureArguments.ManagedAgentCountFn != nil {
 			managedAgentCount, err = featureArguments.ManagedAgentCountFn(ctx, agentLimit.UsagePeriod.Start, agentLimit.UsagePeriod.End)
 		}
-		if err != nil {
+		if xerrors.Is(err, context.Canceled) || xerrors.Is(err, context.DeadlineExceeded) {
+			// If the context is canceled, we want to bail the entire
+			// LicensesEntitlements call.
+			return entitlements, xerrors.Errorf("get managed agent count: %w", err)
+		} else if err != nil {
 			entitlements.Errors = append(entitlements.Errors,
 				fmt.Sprintf("Error getting managed agent count: %s", err.Error()))
 		} else {

Original file line number	Diff line number	Diff line change
`@@ -830,7 +830,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {`
`830`	`830`	`}`
`831`	`831`	`api.derpMesh.SetAddresses(addresses, false)`
`832`	`832`	`}`
`833`		`- _ = api.updateEntitlements(ctx)`
	`833`	`+ _ = api.updateEntitlements(api.ctx)`
`834`	`834`	`})`
`835`	`835`	`} else {`
`836`	`836`	`coordinator = agpltailnet.NewCoordinator(api.Logger)`
`@@ -840,7 +840,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {`
`840`	`840`	`api.replicaManager.SetCallback(func() {`
`841`	`841`	`// If the amount of replicas change, so should our entitlements.`
`842`	`842`	`// This is to display a warning in the UI if the user is unlicensed.`
`843`		`- _ = api.updateEntitlements(ctx)`
	`843`	`+ _ = api.updateEntitlements(api.ctx)`
`844`	`844`	`})`
`845`	`845`	`}`
`846`	`846`