fix: fix subscription check and templating (#5080)

kubeshop · Feb 28, 2024 · 66fb242 · 66fb242
1 parent bc2dc9d
commit 66fb242
Show file tree

Hide file tree

Showing 7 changed files with 117 additions and 50 deletions.
diff --git a/cmd/api-server/main.go b/cmd/api-server/main.go
@@ -488,6 +488,27 @@ func main() {
 		ui.ExitOnError("Creating slack loader", err)
 	}
 
+	proContext := config.ProContext{
+		APIKey:               cfg.TestkubeProAPIKey,
+		URL:                  cfg.TestkubeProURL,
+		LogsPath:             cfg.TestkubeProLogsPath,
+		TLSInsecure:          cfg.TestkubeProTLSInsecure,
+		WorkerCount:          cfg.TestkubeProWorkerCount,
+		LogStreamWorkerCount: cfg.TestkubeProLogStreamWorkerCount,
+		SkipVerify:           cfg.TestkubeProSkipVerify,
+		EnvID:                cfg.TestkubeProEnvID,
+		OrgID:                cfg.TestkubeProOrgID,
+		Migrate:              cfg.TestkubeProMigrate,
+		ConnectionTimeout:    cfg.TestkubeProConnectionTimeout,
+	}
+
+	// Check Pro/Enterprise subscription
+	var subscriptionChecker checktcl.SubscriptionChecker
+	if mode == common.ModeAgent {
+		subscriptionChecker, err = checktcl.NewSubscriptionChecker(ctx, proContext, grpcClient, grpcConn)
+		ui.ExitOnError("Failed creating subscription checker", err)
+	}
+
 	api := apiv1.NewTestkubeAPI(
 		cfg.TestkubeNamespace,
 		resultsRepository,
@@ -521,31 +542,12 @@ func main() {
 		features,
 		logsStream,
 		logGrpcClient,
+		subscriptionChecker,
 	)
 
-	var proContext *config.ProContext
 	if mode == common.ModeAgent {
 		log.DefaultLogger.Info("starting agent service")
-		proContext = &config.ProContext{
-			APIKey:               cfg.TestkubeProAPIKey,
-			URL:                  cfg.TestkubeProURL,
-			LogsPath:             cfg.TestkubeProLogsPath,
-			TLSInsecure:          cfg.TestkubeProTLSInsecure,
-			WorkerCount:          cfg.TestkubeProWorkerCount,
-			LogStreamWorkerCount: cfg.TestkubeProLogStreamWorkerCount,
-			SkipVerify:           cfg.TestkubeProSkipVerify,
-			EnvID:                cfg.TestkubeProEnvID,
-			OrgID:                cfg.TestkubeProOrgID,
-			Migrate:              cfg.TestkubeProMigrate,
-			ConnectionTimeout:    cfg.TestkubeProConnectionTimeout,
-		}
-
-		api.WithProContext(proContext)
-		// Check Pro/Enterprise subscription
-		subscriptionChecker, err := checktcl.NewSubscriptionChecker(ctx, *proContext, grpcClient, grpcConn)
-		ui.WarnOnError("Creating subscription checker", err)
-		api.WithSubscriptionChecker(subscriptionChecker)
-
+		api.WithProContext(&proContext)
 		agentHandle, err := agent.NewAgent(
 			log.DefaultLogger,
 			api.Mux.Handler(),
@@ -555,7 +557,7 @@ func main() {
 			cfg.TestkubeClusterName,
 			envs,
 			features,
-			*proContext,
+			proContext,
 		)
 		if err != nil {
 			ui.ExitOnError("Starting agent", err)
@@ -571,7 +573,6 @@ func main() {
 	}
 
 	api.InitEvents()
-
 	if !cfg.DisableTestTriggers {
 		triggerService := triggers.NewService(
 			sched,

diff --git a/internal/app/api/v1/server.go b/internal/app/api/v1/server.go
@@ -94,6 +94,7 @@ func NewTestkubeAPI(
 	ff featureflags.FeatureFlags,
 	logsStream logsclient.Stream,
 	logGrpcClient logsclient.StreamGetter,
+	subscriptionChecker checktcl.SubscriptionChecker,
 ) TestkubeAPI {
 
 	var httpConfig server.Config
@@ -141,6 +142,7 @@ func NewTestkubeAPI(
 		featureFlags:          ff,
 		logsStream:            logsStream,
 		logGrpcClient:         logGrpcClient,
+		SubscriptionChecker:   subscriptionChecker,
 	}
 
 	// will be reused in websockets handler

diff --git a/pkg/api/v1/testkube/model_test_suite_base_extended.go b/pkg/api/v1/testkube/model_test_suite_base_extended.go
@@ -79,4 +79,52 @@ func (t *TestSuite) QuoteTestSuiteTextFields() {
 			}
 		}
 	}
+	for i := range t.Before {
+		for j := range t.Before[i].Execute {
+			if t.Before[i].Execute[j].ExecutionRequest != nil {
+				t.Before[i].Execute[j].ExecutionRequest.QuoteTestSuiteStepExecutionRequestTextFields()
+			}
+		}
+	}
+	for i := range t.After {
+		for j := range t.After[i].Execute {
+			if t.After[i].Execute[j].ExecutionRequest != nil {
+				t.After[i].Execute[j].ExecutionRequest.QuoteTestSuiteStepExecutionRequestTextFields()
+			}
+		}
+	}
+	for i := range t.Steps {
+		for j := range t.Steps[i].Execute {
+			if t.Steps[i].Execute[j].ExecutionRequest != nil {
+				t.Steps[i].Execute[j].ExecutionRequest.QuoteTestSuiteStepExecutionRequestTextFields()
+			}
+		}
+	}
+}
+
+func (request *TestSuiteStepExecutionRequest) QuoteTestSuiteStepExecutionRequestTextFields() {
+	for i := range request.Args {
+		if request.Args[i] != "" {
+			request.Args[i] = fmt.Sprintf("%q", request.Args[i])
+		}
+	}
+
+	for i := range request.Command {
+		if request.Command[i] != "" {
+			request.Command[i] = fmt.Sprintf("%q", request.Command[i])
+		}
+	}
+
+	var fields = []*string{
+		&request.JobTemplate,
+		&request.CronJobTemplate,
+		&request.ScraperTemplate,
+		&request.PvcTemplate,
+	}
+
+	for _, field := range fields {
+		if *field != "" {
+			*field = fmt.Sprintf("%q", *field)
+		}
+	}
 }
diff --git a/pkg/crd/templates/testsuite.tmpl b/pkg/crd/templates/testsuite.tmpl
@@ -83,7 +83,7 @@ spec:
         {{- if ne (len .ExecutionRequest.Args) 0 }}
         args:
         {{- range .ExecutionRequest.Args }}
-          - "{{ . }}"
+          - {{ . }}
         {{- end }}
         {{- end }}
         {{- if .ExecutionRequest.ArgsMode }}
@@ -214,7 +214,7 @@ spec:
         {{- if ne (len .ExecutionRequest.Args) 0 }}
         args:
         {{- range .ExecutionRequest.Args }}
-          - "{{ . }}"
+          - {{ . }}
         {{- end }}
         {{- end }}
         {{- if .ExecutionRequest.ArgsMode }}
@@ -345,7 +345,7 @@ spec:
         {{- if ne (len .ExecutionRequest.Args) 0 }}
         args:
         {{- range .ExecutionRequest.Args }}
-          - "{{ . }}"
+          - {{ . }}
         {{- end }}
         {{- end }}
         {{- if .ExecutionRequest.ArgsMode }}

diff --git a/pkg/tcl/checktcl/organization_plan.go b/pkg/tcl/checktcl/organization_plan.go
@@ -42,6 +42,22 @@ type OrganizationPlan struct {
 	PlanStatus PlanStatus `json:"planStatus"`
 }
 
+func (p OrganizationPlan) IsEnterprise() bool {
+	return p.TestkubeMode == OrganizationPlanTestkubeModeEnterprise
+}
+
+func (p OrganizationPlan) IsPro() bool {
+	return p.TestkubeMode == OrganizationPlanTestkubeModePro
+}
+
+func (p OrganizationPlan) IsActive() bool {
+	return p.PlanStatus == PlanStatusActive
+}
+
+func (p OrganizationPlan) IsEmpty() bool {
+	return p.PlanStatus == "" && p.TestkubeMode == "" && !p.IsTrial
+}
+
 type GetOrganizationPlanRequest struct{}
 type GetOrganizationPlanResponse struct {
 	TestkubeMode string

diff --git a/pkg/tcl/checktcl/subscription.go b/pkg/tcl/checktcl/subscription.go
@@ -23,7 +23,7 @@ import (
 
 type SubscriptionChecker struct {
 	proContext config.ProContext
-	orgPlan    *OrganizationPlan
+	orgPlan    OrganizationPlan
 }
 
 // NewSubscriptionChecker creates a new subscription checker using the agent token
@@ -47,37 +47,37 @@ func NewSubscriptionChecker(ctx context.Context, proContext config.ProContext, c
 		PlanStatus:   PlanStatus(commandResponse.PlanStatus),
 	}
 
-	return SubscriptionChecker{proContext: proContext, orgPlan: &subscription}, nil
+	return SubscriptionChecker{proContext: proContext, orgPlan: subscription}, nil
 }
 
 // GetCurrentOrganizationPlan returns current organization plan
-func (c *SubscriptionChecker) GetCurrentOrganizationPlan() (*OrganizationPlan, error) {
-	if c.orgPlan == nil {
-		return nil, errors.New("organization plan is not set")
+func (c *SubscriptionChecker) GetCurrentOrganizationPlan() (OrganizationPlan, error) {
+	if c.orgPlan.IsEmpty() {
+		return OrganizationPlan{}, errors.New("organization plan is not set")
 	}
 	return c.orgPlan, nil
 }
 
 // IsOrgPlanEnterprise checks if organization plan is enterprise
 func (c *SubscriptionChecker) IsOrgPlanEnterprise() (bool, error) {
-	if c.orgPlan == nil {
+	if c.orgPlan.IsEmpty() {
 		return false, errors.New("organization plan is not set")
 	}
-	return c.orgPlan.TestkubeMode == OrganizationPlanTestkubeModeEnterprise, nil
+	return c.orgPlan.IsEnterprise(), nil
 }
 
 // IsOrgPlanCloud checks if organization plan is cloud
 func (c *SubscriptionChecker) IsOrgPlanPro() (bool, error) {
-	if c.orgPlan == nil {
+	if c.orgPlan.IsEmpty() {
 		return false, errors.New("organization plan is not set")
 	}
-	return c.orgPlan.TestkubeMode == OrganizationPlanTestkubeModePro, nil
+	return c.orgPlan.IsPro(), nil
 }
 
 // IsOrgPlanActive checks if organization plan is active
 func (c *SubscriptionChecker) IsOrgPlanActive() (bool, error) {
-	if c.orgPlan == nil {
+	if c.orgPlan.IsEmpty() {
 		return false, errors.New("organization plan is not set")
 	}
-	return c.orgPlan.PlanStatus == PlanStatusActive, nil
+	return c.orgPlan.IsActive(), nil
 }
diff --git a/pkg/tcl/checktcl/subscription_test.go b/pkg/tcl/checktcl/subscription_test.go
@@ -16,8 +16,8 @@ import (
 func TestSubscriptionChecker_GetCurrentOrganizationPlan(t *testing.T) {
 	tests := []struct {
 		name    string
-		orgPlan *OrganizationPlan
-		want    *OrganizationPlan
+		orgPlan OrganizationPlan
+		want    OrganizationPlan
 		wantErr bool
 	}{
 		{
@@ -26,12 +26,12 @@ func TestSubscriptionChecker_GetCurrentOrganizationPlan(t *testing.T) {
 		},
 		{
 			name: "Org plan exists",
-			orgPlan: &OrganizationPlan{
+			orgPlan: OrganizationPlan{
 				TestkubeMode: OrganizationPlanTestkubeModeEnterprise,
 				IsTrial:      false,
 				PlanStatus:   PlanStatusActive,
 			},
-			want: &OrganizationPlan{
+			want: OrganizationPlan{
 				TestkubeMode: OrganizationPlanTestkubeModeEnterprise,
 				IsTrial:      false,
 				PlanStatus:   PlanStatusActive,
@@ -59,7 +59,7 @@ func TestSubscriptionChecker_GetCurrentOrganizationPlan(t *testing.T) {
 func TestSubscriptionChecker_IsOrgPlanEnterprise(t *testing.T) {
 	tests := []struct {
 		name    string
-		orgPlan *OrganizationPlan
+		orgPlan OrganizationPlan
 		want    bool
 		wantErr bool
 	}{
@@ -69,15 +69,15 @@ func TestSubscriptionChecker_IsOrgPlanEnterprise(t *testing.T) {
 		},
 		{
 			name: "enterprise org plan",
-			orgPlan: &OrganizationPlan{
+			orgPlan: OrganizationPlan{
 				TestkubeMode: OrganizationPlanTestkubeModeEnterprise,
 			},
 			want:    true,
 			wantErr: false,
 		},
 		{
 			name: "pro org plan",
-			orgPlan: &OrganizationPlan{
+			orgPlan: OrganizationPlan{
 				TestkubeMode: OrganizationPlanTestkubeModePro,
 			},
 			want:    false,
@@ -104,7 +104,7 @@ func TestSubscriptionChecker_IsOrgPlanEnterprise(t *testing.T) {
 func TestSubscriptionChecker_IsOrgPlanPro(t *testing.T) {
 	tests := []struct {
 		name    string
-		orgPlan *OrganizationPlan
+		orgPlan OrganizationPlan
 		want    bool
 		wantErr bool
 	}{
@@ -114,15 +114,15 @@ func TestSubscriptionChecker_IsOrgPlanPro(t *testing.T) {
 		},
 		{
 			name: "enterprise org plan",
-			orgPlan: &OrganizationPlan{
+			orgPlan: OrganizationPlan{
 				TestkubeMode: OrganizationPlanTestkubeModeEnterprise,
 			},
 			want:    false,
 			wantErr: false,
 		},
 		{
 			name: "pro org plan",
-			orgPlan: &OrganizationPlan{
+			orgPlan: OrganizationPlan{
 				TestkubeMode: OrganizationPlanTestkubeModePro,
 			},
 			want:    true,
@@ -149,7 +149,7 @@ func TestSubscriptionChecker_IsOrgPlanPro(t *testing.T) {
 func TestSubscriptionChecker_IsOrgPlanActive(t *testing.T) {
 	tests := []struct {
 		name    string
-		orgPlan *OrganizationPlan
+		orgPlan OrganizationPlan
 		want    bool
 		wantErr bool
 	}{
@@ -159,15 +159,15 @@ func TestSubscriptionChecker_IsOrgPlanActive(t *testing.T) {
 		},
 		{
 			name: "active org plan",
-			orgPlan: &OrganizationPlan{
+			orgPlan: OrganizationPlan{
 				PlanStatus: PlanStatusActive,
 			},
 			want:    true,
 			wantErr: false,
 		},
 		{
 			name: "inactive org plan",
-			orgPlan: &OrganizationPlan{
+			orgPlan: OrganizationPlan{
 				PlanStatus: PlanStatusUnpaid,
 			},
 			want:    false,