[K10-1939] Add new interface Mapper for RegionToZones (#5180)

* Add new interface Mapper for RegionToZones

* Fix EOF

* Fix params

* Change func name

* Address review comments

pkg/blockstorage/awsebs/awsebs.go

@@ -20,9 +20,11 @@ import (
 
 	"github.com/kanisterio/kanister/pkg/blockstorage"
 	ktags "github.com/kanisterio/kanister/pkg/blockstorage/tags"
+	"github.com/kanisterio/kanister/pkg/blockstorage/zone"
 )
 
 var _ blockstorage.Provider = (*ebsStorage)(nil)
+var _ zone.Mapper = (*ebsStorage)(nil)
 
 type ebsStorage struct {
 	ec2Cli *EC2
@@ -411,7 +413,7 @@ func (s *ebsStorage) VolumeCreateFromSnapshot(ctx context.Context, snapshot bloc
 		return nil, errors.Errorf("Required volume fields not available, volumeType: %s, Az: %s, VolumeTags: %v", snapshot.Volume.VolumeType, snapshot.Volume.Az, snapshot.Volume.Tags)
 	}
 
-	z, err := zoneForVolumeCreateFromSnapshot(ctx, snapshot.Region, snapshot.Volume.Az)
+	z, err := zone.FromSourceRegionZone(ctx, s, snapshot.Region, snapshot.Volume.Az)
 	if err != nil {
 		return nil, err
 	}
@@ -589,7 +591,7 @@ func GetRegionFromEC2Metadata() (string, error) {
 	return awsRegion, errors.Wrap(err, "Failed to get AWS Region")
 }
 
-func regionToZones(ctx context.Context, region string) ([]string, error) {
+func (s *ebsStorage) FromRegion(ctx context.Context, region string) ([]string, error) {
 	// Fall back to using a static map.
 	return staticRegionToZones(region)
 }

pkg/blockstorage/awsebs/zone_test.go

@@ -2,8 +2,11 @@ package awsebs
 
 import (
 	"context"
+	"os"
 
 	. "gopkg.in/check.v1"
+
+	"github.com/kanisterio/kanister/pkg/blockstorage/zone"
 )
 
 type ZoneSuite struct{}
@@ -38,12 +41,30 @@ func (s ZoneSuite) TestZoneWithUnknownNodeZones(c *C) {
 			out:    "us-west-2a",
 		},
 	} {
-
-		z, err := zoneWithUnknownNodeZones(ctx, tc.region, tc.in)
+		config := getConfigForTest(c, tc.region)
+		provider, err := NewProvider(config)
+		z, err := zone.WithUnknownNodeZones(ctx, provider.(zone.Mapper), tc.region, tc.in)
 		c.Assert(err, IsNil)
 		c.Assert(z, Not(Equals), "")
 		if tc.out != "" {
 			c.Assert(z, Equals, tc.out)
 		}
 	}
 }
+
+func getConfigForTest(c *C, region string) map[string]string {
+	config := make(map[string]string)
+	config[ConfigRegion] = region
+	accessKey, ok := os.LookupEnv(AccessKeyID)
+	if !ok {
+		c.Skip("The necessary env variable AWS_ACCESS_KEY_ID is not set.")
+	}
+	secretAccessKey, ok := os.LookupEnv(SecretAccessKey)
+	if !ok {
+		c.Skip("The necessary env variable AWS_SECRET_ACCESS_KEY is not set.")
+	}
+	config[AccessKeyID] = accessKey
+	config[SecretAccessKey] = secretAccessKey
+
+	return config
+}

pkg/blockstorage/zone/zone.go

@@ -6,21 +6,73 @@ import (
 	"sort"
 	"strings"
 
-	"github.com/kanisterio/kanister/pkg/kube"
 	"github.com/pkg/errors"
+	log "github.com/sirupsen/logrus"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
+
+	"github.com/kanisterio/kanister/pkg/kube"
+)
+
+type (
+	Mapper interface {
+		FromRegion(ctx context.Context, region string) ([]string, error)
+	}
 )
 
-// GetZoneFromKnownNodeZones get the zone from known node zones
-func GetZoneFromKnownNodeZones(ctx context.Context, sourceZone string, nzs map[string]struct{}) (string, error) {
+// FromSourceRegionZone gets the zones from the given region and sourceZome
+func FromSourceRegionZone(ctx context.Context, m Mapper, region string, sourceZone string) (string, error) {
+	nzs, err := nodeZones(ctx)
+	if err != nil {
+		log.Errorf("Ignoring error getting Node availability zones. Error: %+v", err)
+	}
+	if len(nzs) != 0 {
+		var z string
+		if z, err = getZoneFromKnownNodeZones(ctx, sourceZone, nzs); err == nil && isZoneValid(ctx, m, z, region) {
+			return z, nil
+		}
+	}
+	return WithUnknownNodeZones(ctx, m, region, sourceZone)
+}
+
+func isZoneValid(ctx context.Context, m Mapper, zone, region string) bool {
+	if validZones, err := m.FromRegion(ctx, region); err == nil {
+		for _, z := range validZones {
+			if zone == z {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// WithUnknownNodeZones get the zone list  for the region
+func WithUnknownNodeZones(ctx context.Context, m Mapper, region string, sourceZone string) (string, error) {
+	// We could not the zones of the nodes, so we return an arbitrary one.
+	zs, err := m.FromRegion(ctx, region)
+	if err != nil || len(zs) == 0 {
+		// If all else fails, we return the original AZ.
+		log.Errorf("Using original AZ. region: %s, Error: %+v", region, err)
+		return sourceZone, nil
+	}
+	// We look for a zone with the same suffix.
+	for _, z := range zs {
+		if getZoneSuffixesMatch(z, sourceZone) {
+			return z, nil
+		}
+	}
+	// We return an arbitrary zone in the region.
+	return zs[0], nil
+}
+
+func getZoneFromKnownNodeZones(ctx context.Context, sourceZone string, nzs map[string]struct{}) (string, error) {
 	// If the original zone is available, we return that one.
 	if _, ok := nzs[sourceZone]; ok {
 		return sourceZone, nil
 	}
 	// If there's an available zone with the zone suffix, we use that one.
 	for nz := range nzs {
-		if GetZoneSuffixesMatch(nz, sourceZone) {
+		if getZoneSuffixesMatch(nz, sourceZone) {
 			return nz, nil
 		}
 	}
@@ -47,8 +99,7 @@ func consistentZone(sourceZone string, nzs map[string]struct{}) (string, error)
 	return s[i], nil
 }
 
-// GetZoneSuffixesMatch check if the given zones have a matching suffix
-func GetZoneSuffixesMatch(zone1, zone2 string) bool {
+func getZoneSuffixesMatch(zone1, zone2 string) bool {
 	a1 := zone1[len(zone1)-1]
 	a2 := zone2[len(zone2)-1]
 	return a1 == a2
@@ -58,8 +109,7 @@ const (
 	nodeZonesErr = `Failed to get Node availability zones.`
 )
 
-// NodeZones get the zones available for the nodes
-func NodeZones(ctx context.Context) (map[string]struct{}, error) {
+func nodeZones(ctx context.Context) (map[string]struct{}, error) {
 	cfg, err := kube.LoadConfig()
 	if err != nil {
 		return nil, errors.Wrap(err, nodeZonesErr)

pkg/blockstorage/zone/zone_kube_test.go

@@ -4,20 +4,17 @@ package zone
 
 import (
 	"context"
-	"fmt"
 
 	. "gopkg.in/check.v1"
 )
 
-type KubeTestAWSEBSSuite struct{}
+type KubeTestZoneSuite struct{}
 
-var _ = Suite(&KubeTestAWSEBSSuite{})
+var _ = Suite(&KubeTestZoneSuite{})
 
-func (s KubeTestAWSEBSSuite) TestNodeZones(c *C) {
-	// skipping this test since it fails on travis(minikube)
-	c.Skip(fmt.Sprintf("Skipping TestNodeZones"))
+func (s KubeTestZoneSuite) TestNodeZones(c *C) {
 	ctx := context.Background()
-	zones, err := NodeZones(ctx)
+	zones, err := nodeZones(ctx)
 	c.Assert(err, IsNil)
 	c.Assert(zones, Not(HasLen), 0)
 }