chore: fix TestManagedAgentLimit flake (#19026)

deansheather · web-flow · commit 9a05b4679b9b · 2025-07-24T05:13:15.000Z
Closes coder/internal#812
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,10 +432,15 @@ func LicensesEntitlements(
 		if featureArguments.ManagedAgentCountFn != nil {
 			managedAgentCount, err = featureArguments.ManagedAgentCountFn(ctx, agentLimit.UsagePeriod.Start, agentLimit.UsagePeriod.End)
 		}
-		if err != nil {
+		switch {
+		case 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)
+		case err != nil:
 			entitlements.Errors = append(entitlements.Errors,
 				fmt.Sprintf("Error getting managed agent count: %s", err.Error()))
-		} else {
+		default:
 			agentLimit.Actual = &managedAgentCount
 			entitlements.AddFeature(codersdk.FeatureManagedAgentLimit, agentLimit)
 

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`