refactor: Quota service

internal/dinosaur/pkg/services/dinosaur.go

@@ -212,17 +212,16 @@ func (k *dinosaurService) DetectInstanceType(dinosaurRequest *dbapi.CentralReque
 		return types.EVAL
 	}
 
-	hasQuota, err := quotaService.HasQuotaAllowance(dinosaurRequest, types.STANDARD)
+	hasQuota, err := quotaService.HasQuotaAllowance(dinosaurRequest)
 	if err != nil {
-		glog.Error(errors.NewWithCause(errors.ErrorGeneral, err, "unable to check quota"))
+		glog.Error(errors.NewWithCause(err.Code, err, "unable to check quota"))
 		return types.EVAL
 	}
 	if hasQuota {
-		glog.Infof("Quota detected for central request %s with quota type %s. Granting instance type %s.", dinosaurRequest.ID, quotaType, types.STANDARD)
-		return types.STANDARD
+		glog.Infof("Quota detected for central request %s with quota type %s. Granting instance type %s.", dinosaurRequest.ID, quotaType, dinosaurRequest.InstanceType)
+		return types.DinosaurInstanceType(dinosaurRequest.InstanceType)
 	}
-
-	glog.Infof("No quota detected for central request %s with quota type %s. Granting instance type %s.", dinosaurRequest.ID, quotaType, types.EVAL)
+	glog.Infof("No quota detected for central request %s with quota type %s. Granting instance type %s.", dinosaurRequest.ID, quotaType, dinosaurRequest.InstanceType)
 	return types.EVAL
 }
 
@@ -255,7 +254,7 @@ func (k *dinosaurService) reserveQuota(ctx context.Context, dinosaurRequest *dba
 	if factoryErr != nil {
 		return "", errors.NewWithCause(errors.ErrorGeneral, factoryErr, "unable to check quota")
 	}
-	subscriptionID, err = quotaService.ReserveQuota(ctx, dinosaurRequest, types.DinosaurInstanceType(dinosaurRequest.InstanceType))
+	subscriptionID, err = quotaService.ReserveQuota(ctx, dinosaurRequest)
 	return subscriptionID, err
 }
 
@@ -274,6 +273,7 @@ func (k *dinosaurService) RegisterDinosaurJob(ctx context.Context, dinosaurReque
 		return errors.TooManyDinosaurInstancesReached(errorMsg)
 	}
 
+	dinosaurRequest.InstanceType = types.STANDARD.String()
 	instanceType := k.DetectInstanceType(dinosaurRequest)
 
 	dinosaurRequest.InstanceType = instanceType.String()

internal/dinosaur/pkg/services/dinosaur_test.go

@@ -9,9 +9,11 @@ import (
 	"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/converters"
+	"github.com/stackrox/acs-fleet-manager/internal/dinosaur/pkg/dinosaurs/types"
 	"github.com/stackrox/acs-fleet-manager/pkg/api"
 	"github.com/stackrox/acs-fleet-manager/pkg/auth"
 	"github.com/stackrox/acs-fleet-manager/pkg/db"
+	"github.com/stackrox/acs-fleet-manager/pkg/errors"
 	"github.com/stretchr/testify/assert"
 	"gorm.io/gorm"
 )
@@ -203,3 +205,68 @@ func Test_dinosaurService_DeprovisionExpiredDinosaursQuery(t *testing.T) {
 	assert.Nil(t, svcErr)
 	assert.True(t, m.Triggered)
 }
+
+func Test_dinosaurService_DetectInstanceType(t *testing.T) {
+	mockedQuotaService := &QuotaServiceMock{}
+	svc := &dinosaurService{
+		connectionFactory: db.NewMockConnectionFactory(nil),
+		quotaServiceFactory: &QuotaServiceFactoryMock{
+			GetQuotaServiceFunc: func(_ api.QuotaType) (QuotaService, *errors.ServiceError) {
+				return mockedQuotaService, nil
+			},
+		},
+		dinosaurConfig: &config.CentralConfig{
+			Quota: &config.CentralQuotaConfig{Type: "test"},
+		},
+	}
+	tests := map[string]struct {
+		requestedType types.DinosaurInstanceType
+		hasQuota      bool
+		hasError      *errors.ServiceError
+
+		expectedType types.DinosaurInstanceType
+	}{
+		"has standard quota for standard request": {
+			requestedType: types.STANDARD,
+			hasQuota:      true,
+			expectedType:  types.STANDARD,
+		},
+		"has eval quota for eval request": {
+			requestedType: types.EVAL,
+			hasQuota:      true,
+			expectedType:  types.EVAL,
+		},
+		"has no standard quota for standard request": {
+			requestedType: types.STANDARD,
+			hasQuota:      false,
+			expectedType:  types.EVAL,
+		},
+		"has no eval quota for eval request": {
+			requestedType: types.EVAL,
+			hasQuota:      false,
+			expectedType:  types.EVAL,
+		},
+		"has error for standard request": {
+			requestedType: types.STANDARD,
+			hasError:      errors.FailedToCheckQuota("test"),
+			expectedType:  types.EVAL,
+		},
+		"has error for eval request": {
+			requestedType: types.EVAL,
+			hasError:      errors.FailedToCheckQuota("test"),
+			expectedType:  types.EVAL,
+		},
+	}
+	central := buildCentralRequest(nil)
+	for test, args := range tests {
+		t.Run(test, func(tt *testing.T) {
+			central.InstanceType = args.requestedType.String()
+			mockedQuotaService.HasQuotaAllowanceFunc = func(dinosaur *dbapi.CentralRequest) (bool, *errors.ServiceError) {
+				return args.hasQuota, args.hasError
+			}
+			instanceType := svc.DetectInstanceType(central)
+			assert.Equal(tt, args.requestedType.String(), central.InstanceType, "central must not be changed")
+			assert.Equal(tt, args.expectedType, instanceType)
+		})
+	}
+}

internal/dinosaur/pkg/services/quota.go

@@ -4,7 +4,6 @@ import (
 	"context"
 
 	"github.com/stackrox/acs-fleet-manager/internal/dinosaur/pkg/api/dbapi"
-	"github.com/stackrox/acs-fleet-manager/internal/dinosaur/pkg/dinosaurs/types"
 	"github.com/stackrox/acs-fleet-manager/pkg/errors"
 )
 
@@ -13,9 +12,9 @@ import (
 //go:generate moq -out quotaservice_moq.go . QuotaService
 type QuotaService interface {
 	// HasQuotaAllowance checks if allowed quota is not zero for the given instance type
-	HasQuotaAllowance(dinosaur *dbapi.CentralRequest, instanceType types.DinosaurInstanceType) (bool, *errors.ServiceError)
+	HasQuotaAllowance(dinosaur *dbapi.CentralRequest) (bool, *errors.ServiceError)
 	// ReserveQuota reserves a quota for a user and return the reservation id or an error in case of failure
-	ReserveQuota(ctx context.Context, dinosaur *dbapi.CentralRequest, instanceType types.DinosaurInstanceType) (string, *errors.ServiceError)
+	ReserveQuota(ctx context.Context, dinosaur *dbapi.CentralRequest) (string, *errors.ServiceError)
 	// DeleteQuota deletes a reserved quota
 	DeleteQuota(subscriptionID string) *errors.ServiceError
 }

internal/dinosaur/pkg/services/quota/ams_quota_service.go

@@ -40,13 +40,13 @@ var supportedAMSBillingModels = map[string]struct{}{
 }
 
 // HasQuotaAllowance checks if allowed quota is not zero for the given instance type.
-func (q amsQuotaService) HasQuotaAllowance(central *dbapi.CentralRequest, instanceType types.DinosaurInstanceType) (bool, *errors.ServiceError) {
+func (q amsQuotaService) HasQuotaAllowance(central *dbapi.CentralRequest) (bool, *errors.ServiceError) {
 	org, err := q.amsClient.GetOrganisationFromExternalID(central.OrganisationID)
 	if err != nil {
 		return false, errors.OrganisationNotFound(central.OrganisationID, err)
 	}
 
-	quotaType := instanceType.GetQuotaType()
+	quotaType := types.DinosaurInstanceType(central.InstanceType).GetQuotaType()
 	quotaCosts, err := q.amsClient.GetQuotaCostsForProduct(org.ID(), quotaType.GetResourceName(), quotaType.GetProduct())
 	if err != nil {
 		return false, errors.NewWithCause(errors.ErrorGeneral, err, fmt.Sprintf(
@@ -114,7 +114,8 @@ func (q amsQuotaService) selectBillingModelFromDinosaurInstanceType(orgID, cloud
 }
 
 // ReserveQuota ...
-func (q amsQuotaService) ReserveQuota(ctx context.Context, dinosaur *dbapi.CentralRequest, instanceType types.DinosaurInstanceType) (string, *errors.ServiceError) {
+func (q amsQuotaService) ReserveQuota(ctx context.Context, dinosaur *dbapi.CentralRequest) (string, *errors.ServiceError) {
+	instanceType := types.DinosaurInstanceType(dinosaur.InstanceType)
 	dinosaurID := dinosaur.ID
 	rr := newBaseQuotaReservedResourceResourceBuilder()
 

internal/dinosaur/pkg/services/quota/ams_quota_service_test.go

@@ -198,14 +198,18 @@ func Test_AMSCheckQuota(t *testing.T) {
 				},
 				Owner: tt.args.owner,
 			}
-			standardAllowance, err := quotaService.HasQuotaAllowance(dinosaur, types.STANDARD)
+			dinosaur.InstanceType = types.STANDARD.String()
+			standardAllowance, err := quotaService.HasQuotaAllowance(dinosaur)
 			gomega.Expect(err).ToNot(gomega.HaveOccurred())
-			evalAllowance, err := quotaService.HasQuotaAllowance(dinosaur, types.EVAL)
+
+			dinosaur.InstanceType = types.EVAL.String()
+			evalAllowance, err := quotaService.HasQuotaAllowance(dinosaur)
 			gomega.Expect(err).ToNot(gomega.HaveOccurred())
 			gomega.Expect(standardAllowance).To(gomega.Equal(tt.args.hasStandardQuota))
 			gomega.Expect(evalAllowance).To(gomega.Equal(tt.args.hasEvalQuota))
 
-			_, err = quotaService.ReserveQuota(emptyCtx, dinosaur, tt.args.dinosaurInstanceType)
+			dinosaur.InstanceType = tt.args.dinosaurInstanceType.String()
+			_, err = quotaService.ReserveQuota(emptyCtx, dinosaur)
 			gomega.Expect(err != nil).To(gomega.Equal(tt.wantErr))
 		})
 	}
@@ -660,8 +664,9 @@ func Test_AMSReserveQuota(t *testing.T) {
 				Owner:          tt.args.owner,
 				CloudAccountID: tt.args.cloudAccountID,
 				CloudProvider:  utils.IfThenElse(tt.args.cloudProviderID == "", "cloudProviderID", tt.args.cloudProviderID),
+				InstanceType:   string(types.STANDARD),
 			}
-			subID, err := quotaService.ReserveQuota(emptyCtx, dinosaur, types.STANDARD)
+			subID, err := quotaService.ReserveQuota(emptyCtx, dinosaur)
 			gomega.Expect(subID).To(gomega.Equal(tt.want))
 			gomega.Expect(err != nil).To(gomega.Equal(tt.wantErr))
 
@@ -752,17 +757,13 @@ func Test_Delete_Quota(t *testing.T) {
 }
 
 func Test_amsQuotaService_HasQuotaAllowance(t *testing.T) {
-	type args struct {
-		dinosaurRequest      *dbapi.CentralRequest
-		dinosaurInstanceType types.DinosaurInstanceType
-	}
 
 	tests := []struct {
-		name      string
-		ocmClient ocm.Client
-		args      args
-		want      bool
-		wantErr   bool
+		name            string
+		ocmClient       ocm.Client
+		dinosaurRequest *dbapi.CentralRequest
+		want            bool
+		wantErr         bool
 	}{
 		{
 			name: "returns false if no quota cost exists for the dinosaur's organization",
@@ -775,9 +776,9 @@ func Test_amsQuotaService_HasQuotaAllowance(t *testing.T) {
 					return []*v1.QuotaCost{}, nil
 				},
 			},
-			args: args{
-				dinosaurRequest:      &dbapi.CentralRequest{OrganisationID: "dinosaur-org-1"},
-				dinosaurInstanceType: types.STANDARD,
+			dinosaurRequest: &dbapi.CentralRequest{
+				OrganisationID: "dinosaur-org-1",
+				InstanceType:   types.STANDARD.String(),
 			},
 			want:    false,
 			wantErr: false,
@@ -799,9 +800,10 @@ func Test_amsQuotaService_HasQuotaAllowance(t *testing.T) {
 					return []*v1.QuotaCost{qcb}, nil
 				},
 			},
-			args: args{
-				dinosaurRequest:      &dbapi.CentralRequest{OrganisationID: "dinosaur-org-1"},
-				dinosaurInstanceType: types.STANDARD,
+
+			dinosaurRequest: &dbapi.CentralRequest{
+				OrganisationID: "dinosaur-org-1",
+				InstanceType:   types.STANDARD.String(),
 			},
 			want:    false,
 			wantErr: true,
@@ -823,9 +825,9 @@ func Test_amsQuotaService_HasQuotaAllowance(t *testing.T) {
 					return []*v1.QuotaCost{qcb}, nil
 				},
 			},
-			args: args{
-				dinosaurRequest:      &dbapi.CentralRequest{OrganisationID: "dinosaur-org-1"},
-				dinosaurInstanceType: types.STANDARD,
+			dinosaurRequest: &dbapi.CentralRequest{
+				OrganisationID: "dinosaur-org-1",
+				InstanceType:   types.STANDARD.String(),
 			},
 			want:    true,
 			wantErr: false,
@@ -853,12 +855,11 @@ func Test_amsQuotaService_HasQuotaAllowance(t *testing.T) {
 					return []*v1.QuotaCost{qcb, qcb2}, nil
 				},
 			},
-			args: args{
-				dinosaurRequest: &dbapi.CentralRequest{OrganisationID: "dinosaur-org-1",
-					CloudProvider:  awsCloudProvider,
-					CloudAccountID: "cloudAccountID",
-				},
-				dinosaurInstanceType: types.STANDARD,
+			dinosaurRequest: &dbapi.CentralRequest{
+				OrganisationID: "dinosaur-org-1",
+				CloudProvider:  awsCloudProvider,
+				CloudAccountID: "cloudAccountID",
+				InstanceType:   types.STANDARD.String(),
 			},
 			want:    true,
 			wantErr: false,
@@ -884,9 +885,9 @@ func Test_amsQuotaService_HasQuotaAllowance(t *testing.T) {
 					return []*v1.QuotaCost{qcb, qcb2}, nil
 				},
 			},
-			args: args{
-				dinosaurRequest:      &dbapi.CentralRequest{OrganisationID: "dinosaur-org-1"},
-				dinosaurInstanceType: types.STANDARD,
+			dinosaurRequest: &dbapi.CentralRequest{
+				OrganisationID: "dinosaur-org-1",
+				InstanceType:   types.STANDARD.String(),
 			},
 			want:    false,
 			wantErr: false,
@@ -901,9 +902,9 @@ func Test_amsQuotaService_HasQuotaAllowance(t *testing.T) {
 					return []*v1.QuotaCost{}, nil
 				},
 			},
-			args: args{
-				dinosaurRequest:      &dbapi.CentralRequest{OrganisationID: "dinosaur-org-1"},
-				dinosaurInstanceType: types.STANDARD,
+			dinosaurRequest: &dbapi.CentralRequest{
+				OrganisationID: "dinosaur-org-1",
+				InstanceType:   types.STANDARD.String(),
 			},
 			wantErr: true,
 		},
@@ -918,9 +919,9 @@ func Test_amsQuotaService_HasQuotaAllowance(t *testing.T) {
 					return []*v1.QuotaCost{}, fmt.Errorf("error getting quota costs")
 				},
 			},
-			args: args{
-				dinosaurRequest:      &dbapi.CentralRequest{OrganisationID: "dinosaur-org-1"},
-				dinosaurInstanceType: types.STANDARD,
+			dinosaurRequest: &dbapi.CentralRequest{
+				OrganisationID: "dinosaur-org-1",
+				InstanceType:   types.STANDARD.String(),
 			},
 			wantErr: true,
 		},
@@ -931,7 +932,7 @@ func Test_amsQuotaService_HasQuotaAllowance(t *testing.T) {
 			gomega.RegisterTestingT(t)
 			quotaServiceFactory := NewDefaultQuotaServiceFactory(tt.ocmClient, nil, nil)
 			quotaService, _ := quotaServiceFactory.GetQuotaService(api.AMSQuotaType)
-			res, err := quotaService.HasQuotaAllowance(tt.args.dinosaurRequest, tt.args.dinosaurInstanceType)
+			res, err := quotaService.HasQuotaAllowance(tt.dinosaurRequest)
 			gomega.Expect(err != nil).To(gomega.Equal(tt.wantErr))
 			gomega.Expect(res).To(gomega.Equal(tt.want))
 		})
@@ -1144,7 +1145,7 @@ func Test_amsQuotaService_HasQuotaAllowance_Extra(t *testing.T) {
 			quotaServiceFactory := NewDefaultQuotaServiceFactory(amsClient, nil, nil)
 			quotaService, _ := quotaServiceFactory.GetQuotaService(api.AMSQuotaType)
 
-			got, err := quotaService.HasQuotaAllowance(tt.central, types.DinosaurInstanceType(tt.central.InstanceType))
+			got, err := quotaService.HasQuotaAllowance(tt.central)
 			g.Expect(err != nil).To(gomega.Equal(tt.wantErr))
 			if tt.wantErr {
 				g.Expect(err.Error()).To(gomega.Equal(tt.wantErrMsg), err.Error())

internal/dinosaur/pkg/services/quota/quota_management_list_service.go

@@ -19,7 +19,7 @@ type QuotaManagementListService struct {
 }
 
 // HasQuotaAllowance ...
-func (q QuotaManagementListService) HasQuotaAllowance(dinosaur *dbapi.CentralRequest, instanceType types.DinosaurInstanceType) (bool, *errors.ServiceError) {
+func (q QuotaManagementListService) HasQuotaAllowance(dinosaur *dbapi.CentralRequest) (bool, *errors.ServiceError) {
 	username := dinosaur.Owner
 	orgID := dinosaur.OrganisationID
 	org, orgFound := q.quotaManagementList.QuotaList.Organisations.GetByID(orgID)
@@ -33,7 +33,7 @@ func (q QuotaManagementListService) HasQuotaAllowance(dinosaur *dbapi.CentralReq
 		userIsRegistered = userFound
 		allowed = user.GetMaxAllowedInstances() > 0
 	}
-
+	instanceType := types.DinosaurInstanceType(dinosaur.InstanceType)
 	// allow user defined in quota list to create standard instances, and
 	// allow user who are not in quota list to create eval instances.
 	if userIsRegistered && instanceType == types.STANDARD ||
@@ -45,7 +45,8 @@ func (q QuotaManagementListService) HasQuotaAllowance(dinosaur *dbapi.CentralReq
 }
 
 // ReserveQuota ...
-func (q QuotaManagementListService) ReserveQuota(_ context.Context, dinosaur *dbapi.CentralRequest, instanceType types.DinosaurInstanceType) (string, *errors.ServiceError) {
+func (q QuotaManagementListService) ReserveQuota(_ context.Context, dinosaur *dbapi.CentralRequest) (string, *errors.ServiceError) {
+	instanceType := types.DinosaurInstanceType(dinosaur.InstanceType)
 	if !q.quotaManagementList.EnableInstanceLimitControl {
 		return "", nil
 	}

internal/dinosaur/pkg/services/quota/quota_management_list_service_test.go

@@ -148,8 +148,9 @@ func Test_QuotaManagementListCheckQuota(t *testing.T) {
 			dinosaur := &dbapi.CentralRequest{
 				Owner:          "username",
 				OrganisationID: "org-id",
+				InstanceType:   tt.args.instanceType.String(),
 			}
-			allowed, _ := quotaService.HasQuotaAllowance(dinosaur, tt.args.instanceType)
+			allowed, _ := quotaService.HasQuotaAllowance(dinosaur)
 			gomega.Expect(tt.want).To(gomega.Equal(allowed))
 		})
 	}
@@ -370,8 +371,9 @@ func Test_QuotaManagementListReserveQuota(t *testing.T) {
 			dinosaur := &dbapi.CentralRequest{
 				Owner:          "username",
 				OrganisationID: "org-id",
+				InstanceType:   string(tt.args.instanceType),
 			}
-			_, err := quotaService.ReserveQuota(context.Background(), dinosaur, tt.args.instanceType)
+			_, err := quotaService.ReserveQuota(context.Background(), dinosaur)
 			gomega.Expect(tt.wantErr).To(gomega.Equal(err))
 		})
 	}