Packages: Fix `upload-packages` case for Enterprise2 mode (#61772)

* Add corner case for Enterprise2

* Add versionMode corner case

* Fix lint

* Refactor

* Add tests

* Remove break

* Fix lint

* Return error instead of panicing

* Fix lint

pkg/build/cmd/uploadpackages.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"github.com/grafana/grafana/pkg/build/config"
+	"github.com/grafana/grafana/pkg/build/droneutil"
 	"github.com/grafana/grafana/pkg/build/gcloud"
 	"github.com/grafana/grafana/pkg/build/packaging"
 	"github.com/urfave/cli/v2"
@@ -22,10 +23,11 @@ const releaseBranchFolder = "prerelease"
 type uploadConfig struct {
 	config.Config
 
-	edition     config.Edition
-	versionMode config.VersionMode
-	gcpKey      string
-	distDir     string
+	edition       config.Edition
+	versionMode   config.VersionMode
+	gcpKey        string
+	distDir       string
+	versionFolder string
 }
 
 // UploadPackages implements the sub-command "upload-packages".
@@ -94,14 +96,23 @@ func UploadPackages(c *cli.Context) error {
 		distDir:     distDir,
 	}
 
+	event, err := droneutil.GetDroneEventFromEnv()
+	if err != nil {
+		return err
+	}
+
 	if cfg.edition == config.EditionEnterprise2 {
-		if releaseModeConfig.Buckets.ArtifactsEnterprise2 != "" {
-			cfg.Bucket = releaseModeConfig.Buckets.ArtifactsEnterprise2
-		} else {
-			return fmt.Errorf("enterprise2 bucket var doesn't exist")
+		cfg.Bucket, err = bucketForEnterprise2(releaseModeConfig, event)
+		if err != nil {
+			return err
 		}
 	}
 
+	cfg.versionFolder, err = getVersionFolder(cfg, event)
+	if err != nil {
+		return err
+	}
+
 	if err := uploadPackages(cfg); err != nil {
 		return cli.Exit(err.Error(), 1)
 	}
@@ -110,6 +121,40 @@ func UploadPackages(c *cli.Context) error {
 	return nil
 }
 
+// Corner case for custom enterprise2 mode
+func bucketForEnterprise2(releaseModeConfig *config.BuildConfig, event string) (string, error) {
+	if event == config.Custom {
+		buildConfig, err := config.GetBuildConfig(config.ReleaseBranchMode)
+		if err != nil {
+			return "", err
+		}
+		return buildConfig.Buckets.ArtifactsEnterprise2, nil
+	}
+
+	if releaseModeConfig.Buckets.ArtifactsEnterprise2 != "" {
+		return releaseModeConfig.Buckets.ArtifactsEnterprise2, nil
+	}
+
+	return "", fmt.Errorf("enterprise2 bucket var doesn't exist")
+}
+
+func getVersionFolder(cfg uploadConfig, event string) (string, error) {
+	switch cfg.versionMode {
+	case config.TagMode:
+		return releaseFolder, nil
+	case config.MainMode, config.DownstreamMode:
+		return mainFolder, nil
+	case config.ReleaseBranchMode:
+		return releaseBranchFolder, nil
+	default:
+		// Corner case for custom enterprise2 mode
+		if event == config.Custom && cfg.versionMode == config.Enterprise2Mode {
+			return releaseFolder, nil
+		}
+		return "", fmt.Errorf("unrecognized version mode: %s", cfg.versionMode)
+	}
+}
+
 func uploadPackages(cfg uploadConfig) error {
 	log.Printf("Uploading Grafana packages, version %s, %s edition, %s mode...\n", cfg.Version, cfg.edition,
 		cfg.versionMode)
@@ -146,25 +191,13 @@ func uploadPackages(cfg uploadConfig) error {
 		fpaths = append(fpaths, fpath)
 	}
 
-	var versionFolder string
-	switch cfg.versionMode {
-	case config.TagMode:
-		versionFolder = releaseFolder
-	case config.MainMode, config.DownstreamMode:
-		versionFolder = mainFolder
-	case config.ReleaseBranchMode:
-		versionFolder = releaseBranchFolder
-	default:
-		panic(fmt.Sprintf("Unrecognized version mode: %s", cfg.versionMode))
-	}
-
 	var tag, gcsPath string
 	droneTag := strings.TrimSpace(os.Getenv("DRONE_TAG"))
 	if droneTag != "" {
 		tag = droneTag
-		gcsPath = fmt.Sprintf("gs://%s/%s/%s/%s", cfg.Bucket, tag, edition, versionFolder)
+		gcsPath = fmt.Sprintf("gs://%s/%s/%s/%s", cfg.Bucket, tag, edition, cfg.versionFolder)
 	} else {
-		gcsPath = fmt.Sprintf("gs://%s/%s/%s/", cfg.Bucket, edition, versionFolder)
+		gcsPath = fmt.Sprintf("gs://%s/%s/%s/", cfg.Bucket, edition, cfg.versionFolder)
 	}
 	log.Printf("Uploading %d file(s) to GCS (%s)...\n", len(fpaths), gcsPath)
 

pkg/build/cmd/uploadpackages_test.go

@@ -0,0 +1,66 @@
+package main
+
+import (
+	"errors"
+	"fmt"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/build/config"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_getVersionFolder(t *testing.T) {
+	type args struct {
+		cfg           uploadConfig
+		event         string
+		versionFolder string
+	}
+	tests := []struct {
+		name string
+		args args
+		err  error
+	}{
+		{"tag mode", args{uploadConfig{versionMode: config.TagMode}, "", releaseFolder}, nil},
+		{"main mode", args{uploadConfig{versionMode: config.MainMode}, "", mainFolder}, nil},
+		{"downstream mode", args{uploadConfig{versionMode: config.DownstreamMode}, "", mainFolder}, nil},
+		{"release branch mode", args{uploadConfig{versionMode: config.ReleaseBranchMode}, "", releaseBranchFolder}, nil},
+		{"enterprise pro mode", args{uploadConfig{versionMode: config.Enterprise2Mode}, config.Custom, releaseFolder}, nil},
+		{"unrecognised version mode", args{uploadConfig{versionMode: "foo"}, config.Custom, ""}, errors.New("")},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			versionMode, err := getVersionFolder(tt.args.cfg, tt.args.event)
+			if tt.err != nil {
+				require.Error(t, err)
+			}
+			require.Equal(t, versionMode, tt.args.versionFolder)
+		})
+	}
+}
+
+func Test_checkForEnterprise2Edition(t *testing.T) {
+	type args struct {
+		releaseModeConfig *config.BuildConfig
+		event             string
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+		err  error
+	}{
+		{"event is not custom", args{releaseModeConfig: &config.BuildConfig{Buckets: config.Buckets{ArtifactsEnterprise2: "dummy"}}}, "dummy", nil},
+		{"event is not custom and string is empty", args{releaseModeConfig: &config.BuildConfig{Buckets: config.Buckets{ArtifactsEnterprise2: ""}}}, "", fmt.Errorf("enterprise2 bucket var doesn't exist")},
+		{"event is custom", args{releaseModeConfig: nil, event: "custom"}, "grafana-downloads-enterprise2", nil},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := bucketForEnterprise2(tt.args.releaseModeConfig, tt.args.event)
+			if tt.err != nil {
+				require.Error(t, err)
+			}
+			assert.Equalf(t, tt.want, got, "bucketForEnterprise2(%v, %v)", tt.args.releaseModeConfig, tt.args.event)
+		})
+	}
+}