add unit test for getDefaultVolumeSnapshotLocations

Signed-off-by: Wayne Witzel III <wayne@riotousliving.com>

pkg/cmd/server/server.go

@@ -87,8 +87,9 @@ type serverConfig struct {
 
 func NewCommand() *cobra.Command {
 	var (
-		logLevelFlag = logging.LogLevelFlag(logrus.InfoLevel)
-		config       = serverConfig{
+		volumeSnapshotLocations = flag.NewMap().WithKeyValueDelimiter(":")
+		logLevelFlag            = logging.LogLevelFlag(logrus.InfoLevel)
+		config                  = serverConfig{
 			pluginDir:                      "/plugins",
 			metricsAddress:                 defaultMetricsAddress,
 			defaultBackupLocation:          "default",
@@ -131,6 +132,10 @@ func NewCommand() *cobra.Command {
 			}
 			namespace := getServerNamespace(namespaceFlag)
 
+			if volumeSnapshotLocations.Data() != nil {
+				config.defaultVolumeSnapshotLocations = volumeSnapshotLocations.Data()
+			}
+
 			s, err := newServer(namespace, fmt.Sprintf("%s-%s", c.Parent().Name(), c.Name()), config, logger)
 			cmd.CheckError(err)
 
@@ -146,12 +151,7 @@ func NewCommand() *cobra.Command {
 	command.Flags().BoolVar(&config.restoreOnly, "restore-only", config.restoreOnly, "run in a mode where only restores are allowed; backups, schedules, and garbage-collection are all disabled")
 	command.Flags().StringSliceVar(&config.restoreResourcePriorities, "restore-resource-priorities", config.restoreResourcePriorities, "desired order of resource restores; any resource not in the list will be restored alphabetically after the prioritized resources")
 	command.Flags().StringVar(&config.defaultBackupLocation, "default-backup-storage-location", config.defaultBackupLocation, "name of the default backup storage location")
-
-	volumeSnapshotLocations := flag.NewMap().WithKeyValueDelimiter(":")
-	command.Flags().Var(&volumeSnapshotLocations, "default-volume-snapshot-locations", "list of unique volume providers and default volume snapshot location (provider1:location-01, provider2:location-02, ...)")
-	if volumeSnapshotLocations.Data() != nil {
-		config.defaultVolumeSnapshotLocations = volumeSnapshotLocations.Data()
-	}
+	command.Flags().Var(&volumeSnapshotLocations, "default-volume-snapshot-locations", "list of unique volume providers and default volume snapshot location (provider1:location-01,provider2:location-02,...)")
 
 	return command
 }
@@ -286,7 +286,7 @@ func (s *server) run() error {
 			Warnf("Default backup storage location %q not found; backups must explicitly specify a location", s.config.defaultBackupLocation)
 	}
 
-	defaultVolumeSnapshotLocations, err := s.getDefaultVolumeSnapshotLocations()
+	defaultVolumeSnapshotLocations, err := getDefaultVolumeSnapshotLocations(s.arkClient, s.namespace, s.config.defaultVolumeSnapshotLocations)
 	if err != nil {
 		return err
 	}
@@ -313,24 +313,25 @@ func (s *server) run() error {
 	return nil
 }
 
-func (s *server) getDefaultVolumeSnapshotLocations() (map[string]*api.VolumeSnapshotLocation, error) {
+func getDefaultVolumeSnapshotLocations(arkClient clientset.Interface, namespace string, defaultVolumeSnapshotLocations map[string]string) (map[string]*api.VolumeSnapshotLocation, error) {
 	providerDefaults := make(map[string]*api.VolumeSnapshotLocation)
-	if len(s.config.defaultVolumeSnapshotLocations) == 0 {
+	if len(defaultVolumeSnapshotLocations) == 0 {
 		return providerDefaults, nil
 	}
 
-	volumeSnapshotLocations, err := s.arkClient.ArkV1().VolumeSnapshotLocations(s.namespace).List(metav1.ListOptions{})
+	volumeSnapshotLocations, err := arkClient.ArkV1().VolumeSnapshotLocations(namespace).List(metav1.ListOptions{})
 	if err != nil {
 		return providerDefaults, errors.WithStack(err)
 	}
 
 	providerLocations := make(map[string][]*api.VolumeSnapshotLocation)
-	for _, vsl := range volumeSnapshotLocations.Items {
-		providerLocations[vsl.Spec.Provider] = append(providerLocations[vsl.Spec.Provider], &vsl)
+	for i, vsl := range volumeSnapshotLocations.Items {
+		locations := providerLocations[vsl.Spec.Provider]
+		providerLocations[vsl.Spec.Provider] = append(locations, &volumeSnapshotLocations.Items[i])
 	}
 
 	for provider, locations := range providerLocations {
-		defaultLocation, ok := s.config.defaultVolumeSnapshotLocations[provider]
+		defaultLocation, ok := defaultVolumeSnapshotLocations[provider]
 		if !ok {
 			return providerDefaults, errors.Errorf("missing provider %s. When using default volume snapshot locations, one must exist for every known provider.", provider)
 		}

pkg/cmd/server/server_test.go

@@ -25,6 +25,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/heptio/ark/pkg/apis/ark/v1"
+	fakeclientset "github.com/heptio/ark/pkg/generated/clientset/versioned/fake"
 	arktest "github.com/heptio/ark/pkg/util/test"
 )
 
@@ -97,3 +98,59 @@ func TestArkResourcesExist(t *testing.T) {
 	arkAPIResourceList.APIResources = arkAPIResourceList.APIResources[:3]
 	assert.Error(t, server.arkResourcesExist())
 }
+
+func TestDefaultVolumeSnapshotLocations(t *testing.T) {
+	namespace := "heptio-ark"
+	arkClient := fakeclientset.NewSimpleClientset()
+
+	location := &v1.VolumeSnapshotLocation{ObjectMeta: metav1.ObjectMeta{Name: "location1"}, Spec: v1.VolumeSnapshotLocationSpec{Provider: "provider1"}}
+	arkClient.ArkV1().VolumeSnapshotLocations(namespace).Create(location)
+
+	defaultVolumeSnapshotLocations := make(map[string]string)
+
+	// No defaults
+	volumeSnapshotLocations, err := getDefaultVolumeSnapshotLocations(arkClient, namespace, defaultVolumeSnapshotLocations)
+	assert.Equal(t, 0, len(volumeSnapshotLocations))
+	assert.NoError(t, err)
+
+	// Bad location
+	defaultVolumeSnapshotLocations["provider1"] = "badlocation"
+	volumeSnapshotLocations, err = getDefaultVolumeSnapshotLocations(arkClient, namespace, defaultVolumeSnapshotLocations)
+	assert.Equal(t, 0, len(volumeSnapshotLocations))
+	assert.Error(t, err)
+
+	// Bad provider
+	defaultVolumeSnapshotLocations["provider2"] = "badlocation"
+	volumeSnapshotLocations, err = getDefaultVolumeSnapshotLocations(arkClient, namespace, defaultVolumeSnapshotLocations)
+	assert.Equal(t, 0, len(volumeSnapshotLocations))
+	assert.Error(t, err)
+
+	// Good provider, good location
+	delete(defaultVolumeSnapshotLocations, "provider2")
+	defaultVolumeSnapshotLocations["provider1"] = "location1"
+	volumeSnapshotLocations, err = getDefaultVolumeSnapshotLocations(arkClient, namespace, defaultVolumeSnapshotLocations)
+	assert.Equal(t, 1, len(volumeSnapshotLocations))
+	assert.NoError(t, err)
+
+	location2 := &v1.VolumeSnapshotLocation{ObjectMeta: metav1.ObjectMeta{Name: "location2"}, Spec: v1.VolumeSnapshotLocationSpec{Provider: "provider2"}}
+	arkClient.ArkV1().VolumeSnapshotLocations(namespace).Create(location2)
+
+	// Mutliple Provider/Location 1 good, 1 bad
+	defaultVolumeSnapshotLocations["provider2"] = "badlocation"
+	volumeSnapshotLocations, err = getDefaultVolumeSnapshotLocations(arkClient, namespace, defaultVolumeSnapshotLocations)
+	assert.Error(t, err)
+
+	location21 := &v1.VolumeSnapshotLocation{ObjectMeta: metav1.ObjectMeta{Name: "location2-1"}, Spec: v1.VolumeSnapshotLocationSpec{Provider: "provider2"}}
+	arkClient.ArkV1().VolumeSnapshotLocations(namespace).Create(location21)
+
+	location11 := &v1.VolumeSnapshotLocation{ObjectMeta: metav1.ObjectMeta{Name: "location1-1"}, Spec: v1.VolumeSnapshotLocationSpec{Provider: "provider1"}}
+	arkClient.ArkV1().VolumeSnapshotLocations(namespace).Create(location11)
+
+	// Mutliple Provider/Location all good
+	defaultVolumeSnapshotLocations["provider2"] = "location2"
+	volumeSnapshotLocations, err = getDefaultVolumeSnapshotLocations(arkClient, namespace, defaultVolumeSnapshotLocations)
+	assert.Equal(t, 2, len(volumeSnapshotLocations))
+	assert.NoError(t, err)
+	assert.Equal(t, volumeSnapshotLocations["provider1"].ObjectMeta.Name, "location1")
+	assert.Equal(t, volumeSnapshotLocations["provider2"].ObjectMeta.Name, "location2")
+}