grace period mgr test

internal/dinosaur/pkg/workers/dinosaurmgrs/grace_period_mgr.go

@@ -52,9 +52,9 @@ func (k *GracePeriodManager) Reconcile() []error {
 	glog.Infoln("reconciling centrals")
 	var encounteredErrors []error
 
-	centrals, err := k.dinosaurService.ListByStatus(constants.GetDeletingStatuses()...)
-	if err != nil {
-		return append(encounteredErrors, err)
+	centrals, svcErr := k.dinosaurService.ListByStatus(constants.GetDeletingStatuses()...)
+	if svcErr != nil {
+		return append(encounteredErrors, svcErr)
 	}
 
 	// reconciles grace_from field for central instances
@@ -79,6 +79,12 @@ func contains[T comparable](values []T, s T) bool {
 func (k *GracePeriodManager) reconcileCentralGraceFrom(centrals dbapi.CentralList) serviceErr.ErrorList {
 	glog.Infof("reconciling grace period start date for central instances")
 	var svcErrors serviceErr.ErrorList
+
+	quotaService, factoryErr := k.quotaServiceFactory.GetQuotaService(api.QuotaType(k.dinosaurConfig.Quota.Type))
+	if factoryErr != nil {
+		return append(svcErrors, factoryErr)
+	}
+
 	subscriptionStatusByOrg := map[string]bool{}
 
 	for _, central := range centrals {
@@ -93,12 +99,6 @@ func (k *GracePeriodManager) reconcileCentralGraceFrom(centrals dbapi.CentralLis
 		glog.Infof("checking quota entitlement status for Central instance %q", central.ID)
 		active, exists := subscriptionStatusByOrg[central.OrganisationID]
 		if !exists {
-			quotaService, factoryErr := k.quotaServiceFactory.GetQuotaService(api.QuotaType(k.dinosaurConfig.Quota.Type))
-			if factoryErr != nil {
-				svcErrors = append(svcErrors, factoryErr)
-				continue
-			}
-
 			isActive, err := quotaService.IsQuotaEntitlementActive(central)
 			if err != nil {
 				svcErrors = append(svcErrors, errors.Wrapf(err, "failed to get quota entitlement status of central instance %q", central.ID))
@@ -117,22 +117,20 @@ func (k *GracePeriodManager) reconcileCentralGraceFrom(centrals dbapi.CentralLis
 }
 
 // Updates grace_from field of the given Central instance based on the user/organisation's quota entitlement status
-func (k *GracePeriodManager) updateGraceFromBasedOnQuotaEntitlement(central *dbapi.CentralRequest, isQuotaEntitlementActive bool) error {
+func (k *GracePeriodManager) updateGraceFromBasedOnQuotaEntitlement(central *dbapi.CentralRequest, isQuotaEntitlementActive bool) *serviceErr.ServiceError {
 	// if quota entitlement is active, ensure grace_from is set to null
 	if isQuotaEntitlementActive && central.GraceFrom != nil {
+		central.GraceFrom = nil
 		glog.Infof("updating grace start date of central instance %q to NULL", central.ID)
-		return k.dinosaurService.Updates(central, map[string]interface{}{
-			"grace_from": nil,
-		})
+		return k.dinosaurService.Update(central)
 	}
 
 	// if quota entitlement is not active and grace_from is not already set, set its value based on the current time and grace period allowance
 	if !isQuotaEntitlementActive && central.GraceFrom == nil {
-		graceFromTime := time.Now()
-		glog.Infof("quota entitlement for central instance %q is no longer active, updating grace_from to %q", central.ID, graceFromTime.Format(time.RFC1123Z))
-		return k.dinosaurService.Updates(central, map[string]interface{}{
-			"grace_from": &graceFromTime,
-		})
+		now := time.Now()
+		central.GraceFrom = &now
+		glog.Infof("quota entitlement for central instance %q is no longer active, updating grace_from to %q", central.ID, now.Format(time.RFC1123Z))
+		return k.dinosaurService.Update(central)
 	}
 
 	glog.Infof("no grace_from changes needed for central %q, skipping update", central.ID)

internal/dinosaur/pkg/workers/dinosaurmgrs/grace_period_mgr_test.go

@@ -0,0 +1,121 @@
+package dinosaurmgrs
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/stackrox/acs-fleet-manager/internal/dinosaur/constants"
+	"github.com/stackrox/acs-fleet-manager/internal/dinosaur/pkg/api/dbapi"
+	"github.com/stackrox/acs-fleet-manager/internal/dinosaur/pkg/config"
+	"github.com/stackrox/acs-fleet-manager/internal/dinosaur/pkg/services"
+	"github.com/stackrox/acs-fleet-manager/pkg/api"
+	"github.com/stackrox/acs-fleet-manager/pkg/errors"
+)
+
+func TestGracePeriodManager(t *testing.T) {
+	withEntitlement := func(e bool) (*services.QuotaServiceMock, *services.QuotaServiceFactoryMock) {
+		qs := &services.QuotaServiceMock{
+			IsQuotaEntitlementActiveFunc: func(dinosaur *dbapi.CentralRequest) (bool, error) {
+				return e, nil
+			},
+		}
+		return qs, &services.QuotaServiceFactoryMock{
+			GetQuotaServiceFunc: func(quotaType api.QuotaType) (services.QuotaService, *errors.ServiceError) {
+				return qs, nil
+			},
+		}
+	}
+	withCentrals := func(centrals ...*dbapi.CentralRequest) *services.DinosaurServiceMock {
+		return &services.DinosaurServiceMock{
+			ListByStatusFunc: func(status ...constants.CentralStatus) ([]*dbapi.CentralRequest, *errors.ServiceError) {
+				return centrals, nil
+			},
+			UpdateFunc: func(dinosaurRequest *dbapi.CentralRequest) *errors.ServiceError {
+				return nil
+			},
+		}
+	}
+	defaultCfg := &config.CentralConfig{
+		Quota: config.NewCentralQuotaConfig(),
+	}
+
+	t.Run("no centrals, no problem", func(t *testing.T) {
+		dinoSvc := withCentrals()
+		quotaSvc, quotaFactory := withEntitlement(true)
+		gpm := NewGracePeriodManager(dinoSvc, quotaFactory, defaultCfg)
+		errs := gpm.Reconcile()
+		require.Empty(t, errs)
+		assert.Len(t, dinoSvc.ListByStatusCalls(), 1)
+		assert.Empty(t, dinoSvc.UpdateCalls())
+		assert.Empty(t, quotaSvc.IsQuotaEntitlementActiveCalls())
+		assert.Len(t, quotaFactory.GetQuotaServiceCalls(), 1)
+	})
+
+	t.Run("unset grace", func(t *testing.T) {
+		now := time.Now()
+		central := &dbapi.CentralRequest{GraceFrom: &now}
+		dinoSvc := withCentrals(central)
+		quotaSvc, quotaFactory := withEntitlement(true)
+		gpm := NewGracePeriodManager(dinoSvc, quotaFactory, defaultCfg)
+		errs := gpm.Reconcile()
+		require.Empty(t, errs)
+		assert.Nil(t, central.GraceFrom)
+		assert.Len(t, dinoSvc.ListByStatusCalls(), 1)
+		assert.Len(t, quotaSvc.IsQuotaEntitlementActiveCalls(), 1)
+		assert.Len(t, dinoSvc.UpdateCalls(), 1)
+		assert.Len(t, quotaFactory.GetQuotaServiceCalls(), 1)
+	})
+
+	t.Run("set grace", func(t *testing.T) {
+		now := time.Now()
+		central := &dbapi.CentralRequest{}
+		dinoSvc := withCentrals(central)
+		quotaSvc, quotaFactory := withEntitlement(false)
+		gpm := NewGracePeriodManager(dinoSvc, quotaFactory, defaultCfg)
+		errs := gpm.Reconcile()
+		require.Empty(t, errs)
+		require.NotNil(t, central.GraceFrom)
+		assert.Less(t, now, *central.GraceFrom)
+		assert.Len(t, dinoSvc.ListByStatusCalls(), 1)
+		assert.Len(t, quotaSvc.IsQuotaEntitlementActiveCalls(), 1)
+		assert.Len(t, dinoSvc.UpdateCalls(), 1)
+		assert.Len(t, quotaFactory.GetQuotaServiceCalls(), 1)
+	})
+
+	t.Run("quota cost cache in use", func(t *testing.T) {
+		now := time.Now()
+		centralA := &dbapi.CentralRequest{GraceFrom: &now, OrganisationID: "one"}
+		centralB := &dbapi.CentralRequest{GraceFrom: &now, OrganisationID: "one"}
+		centralC := &dbapi.CentralRequest{GraceFrom: &now, OrganisationID: "another"}
+		dinoSvc := withCentrals(centralA, centralB, centralC)
+		quotaSvc, quotaFactory := withEntitlement(true)
+		gpm := NewGracePeriodManager(dinoSvc, quotaFactory, defaultCfg)
+		errs := gpm.Reconcile()
+		require.Empty(t, errs)
+		assert.Nil(t, centralA.GraceFrom)
+		assert.Nil(t, centralB.GraceFrom)
+		assert.Len(t, dinoSvc.ListByStatusCalls(), 1)
+		assert.Len(t, quotaSvc.IsQuotaEntitlementActiveCalls(), 2)
+		assert.Len(t, dinoSvc.UpdateCalls(), 3)
+		assert.Len(t, quotaFactory.GetQuotaServiceCalls(), 1)
+	})
+
+	t.Run("skip deleting and deprovision statuses", func(t *testing.T) {
+		centralA := &dbapi.CentralRequest{Status: constants.CentralRequestStatusDeleting.String()}
+		centralB := &dbapi.CentralRequest{Status: constants.CentralRequestStatusDeprovision.String()}
+		dinoSvc := withCentrals(centralA, centralB)
+		quotaSvc, quotaFactory := withEntitlement(false)
+		gpm := NewGracePeriodManager(dinoSvc, quotaFactory, defaultCfg)
+		errs := gpm.Reconcile()
+		require.Empty(t, errs)
+		assert.Nil(t, centralA.GraceFrom)
+		assert.Nil(t, centralB.GraceFrom)
+		assert.Len(t, dinoSvc.ListByStatusCalls(), 1)
+		assert.Empty(t, quotaSvc.IsQuotaEntitlementActiveCalls())
+		assert.Empty(t, dinoSvc.UpdateCalls())
+		assert.Len(t, quotaFactory.GetQuotaServiceCalls(), 1)
+	})
+}