fix: refresh entitlements after creating first user (#12285)

coderd/coderd.go

@@ -125,6 +125,8 @@ type Options struct {
 	ExternalAuthConfigs            []*externalauth.Config
 	RealIPConfig                   *httpmw.RealIPConfig
 	TrialGenerator                 func(ctx context.Context, body codersdk.LicensorTrialRequest) error
+	// RefreshEntitlements is used to set correct entitlements after creating first user and generating trial license.
+	RefreshEntitlements func(ctx context.Context) error
 	// TLSCertificates is used to mesh DERP servers securely.
 	TLSCertificates    []tls.Certificate
 	TailnetCoordinator tailnet.Coordinator

coderd/coderdtest/coderdtest.go

@@ -107,6 +107,7 @@ type Options struct {
 	TLSCertificates       []tls.Certificate
 	ExternalAuthConfigs   []*externalauth.Config
 	TrialGenerator        func(ctx context.Context, body codersdk.LicensorTrialRequest) error
+	RefreshEntitlements   func(ctx context.Context) error
 	TemplateScheduleStore schedule.TemplateScheduleStore
 	Coordinator           tailnet.Coordinator
 
@@ -434,6 +435,7 @@ func NewOptions(t testing.TB, options *Options) (func(http.Handler), context.Can
 			AccessControlStore:                 accessControlStore,
 			TLSCertificates:                    options.TLSCertificates,
 			TrialGenerator:                     options.TrialGenerator,
+			RefreshEntitlements:                options.RefreshEntitlements,
 			TailnetCoordinator:                 options.Coordinator,
 			BaseDERPMap:                        derpMap,
 			DERPMapUpdateFrequency:             150 * time.Millisecond,

coderd/users.go

@@ -191,6 +191,16 @@ func (api *API) postFirstUser(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	if api.RefreshEntitlements != nil {
+		err = api.RefreshEntitlements(ctx)
+		if err != nil {
+			api.Logger.Error(ctx, "failed to refresh entitlements after generating trial license")
+			return
+		}
+	} else {
+		api.Logger.Debug(ctx, "entitlements will not be refreshed")
+	}
+
 	telemetryUser := telemetry.ConvertUser(user)
 	// Send the initial users email address!
 	telemetryUser.Email = &user.Email

coderd/users_test.go

@@ -75,10 +75,16 @@ func TestFirstUser(t *testing.T) {
 
 	t.Run("Trial", func(t *testing.T) {
 		t.Parallel()
-		called := make(chan struct{})
+		trialGenerated := make(chan struct{})
+		entitlementsRefreshed := make(chan struct{})
+
 		client := coderdtest.New(t, &coderdtest.Options{
 			TrialGenerator: func(context.Context, codersdk.LicensorTrialRequest) error {
-				close(called)
+				close(trialGenerated)
+				return nil
+			},
+			RefreshEntitlements: func(context.Context) error {
+				close(entitlementsRefreshed)
 				return nil
 			},
 		})
@@ -94,7 +100,9 @@ func TestFirstUser(t *testing.T) {
 		}
 		_, err := client.CreateFirstUser(ctx, req)
 		require.NoError(t, err)
-		<-called
+
+		_ = testutil.RequireRecvCtx(ctx, t, trialGenerated)
+		_ = testutil.RequireRecvCtx(ctx, t, entitlementsRefreshed)
 	})
 }