mock_client.go

package storageos

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/hashicorp/go-multierror"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"

	stosv1 "github.com/storageos/api-manager/api/v1"
	stosapi "github.com/storageos/go-api/autogenerated"
	"k8s.io/apimachinery/pkg/api/resource"
)

func init() {
	rand.Seed(int64(time.Now().Nanosecond()))
}

// MockObject can be be used to replace an api object.
type MockObject struct {
	ID        string
	Name      string
	Namespace string
	Labels    map[string]string
	Healthy   bool
	Inner     interface{}
}

func (m MockObject) GetInner() interface{} {
	return m.Inner
}

// GetID returns the object ID.
func (m MockObject) GetID() string {
	return m.ID
}

// GetName returns the object name.
func (m MockObject) GetName() string {
	return m.Name
}

// GetNamespace returns the object namespace.
func (m MockObject) GetNamespace() string {
	return m.Namespace
}

// GetLabels returns the object labels.
func (m MockObject) GetLabels() map[string]string {
	return m.Labels
}

// IsHealthy returns true if the object is healthy.
func (m MockObject) IsHealthy() bool {
	return m.Healthy
}

// MockLicence is used to replace a licence.
type MockLicence struct {
	ClusterID            string
	ExpiresAt            time.Time
	ClusterCapacityBytes uint64
	UsedBytes            uint64
	Kind                 string
	CustomerName         string
	Features             *[]string
	Version              string
}

func (l MockLicence) GetClusterID() string {
	return l.ClusterID
}

func (l MockLicence) GetExpiresAt() time.Time {
	return l.ExpiresAt
}

func (l MockLicence) GetClusterCapacityBytes() uint64 {
	return l.ClusterCapacityBytes
}

func (l MockLicence) GetUsedBytes() uint64 {
	return l.UsedBytes
}

func (l MockLicence) GetKind() string {
	return l.Kind
}

func (l MockLicence) GetCustomerName() string {
	return l.CustomerName
}

func (l MockLicence) GetFeatures() *[]string {
	return l.Features
}

func (l MockLicence) GetVersion() string {
	return l.Version
}

// MockClient provides a test interface to the StorageOS api.
type MockClient struct {
	setExternalEndpoint      func(string, string) error
	namespaces               map[client.ObjectKey]Object
	nodes                    map[client.ObjectKey]Object
	volumes                  map[client.ObjectKey]Object
	nodeLabels               map[string]string
	licence                  LicenceObject
	mu                       sync.RWMutex
	DeleteNamespaceCallCount map[client.ObjectKey]int
	DeleteNodeCallCount      map[client.ObjectKey]int
	ListNamespacesErr        error
	DeleteNamespaceErr       error
	GetNodeErr               error
	NodeObjectsErr           error
	ListNodesErr             error
	DeleteNodeErr            error
	EnsureNodeLabelsErr      error
	GetNodeLabelsErr         error
	ListVolumeErr            error
	GetVolumeErr             error
	VolumeObjectsErr         error
	EnsureVolumeLabelsErr    error
	SharedVolsErr            error
	SharedVolErr             error
	SetEndpointErr           error
	GetLicenceErr            error
}

// NewMockClient returns an initialized MockClient.
func NewMockClient() *MockClient {
	var licenceObject LicenceObject
	return &MockClient{
		namespaces:               make(map[client.ObjectKey]Object),
		nodes:                    make(map[client.ObjectKey]Object),
		volumes:                  make(map[client.ObjectKey]Object),
		nodeLabels:               make(map[string]string),
		licence:                  licenceObject,
		DeleteNamespaceCallCount: make(map[client.ObjectKey]int),
		DeleteNodeCallCount:      make(map[client.ObjectKey]int),
		mu:                       sync.RWMutex{},
	}
}

// ListNamespaces returns a list of StorageOS namespace objects.
func (c *MockClient) ListNamespaces(ctx context.Context) ([]Object, error) {
	if c.ListNamespacesErr != nil {
		return nil, c.ListNamespacesErr
	}
	ret := []Object{}
	c.mu.RLock()
	defer c.mu.RUnlock()
	for _, ns := range c.namespaces {
		ret = append(ret, ns)
	}
	return ret, nil
}

// AddNamespace adds a namespace to the StorageOS cluster.
func (c *MockClient) AddNamespace(key client.ObjectKey) error {
	c.mu.Lock()
	c.namespaces[key] = MockObject{Name: key.Name}
	c.mu.Unlock()
	return nil
}

// NamespaceExists returns true if the namespace exists.
func (c *MockClient) NamespaceExists(key client.ObjectKey) bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	if _, ok := c.namespaces[key]; ok {
		return true
	}
	return false
}

// DeleteNamespace removes a namespace from the StorageOS cluster.
func (c *MockClient) DeleteNamespace(ctx context.Context, key client.ObjectKey) error {
	c.DeleteNamespaceCallCount[key]++
	if c.DeleteNamespaceErr != nil {
		return c.DeleteNamespaceErr
	}
	if !c.NamespaceExists(key) {
		return ErrNamespaceNotFound
	}
	c.mu.Lock()
	delete(c.namespaces, key)
	c.mu.Unlock()
	return nil
}

// NodeObjects returns a map of nodes objects, indexed on object key.
func (c *MockClient) NodeObjects(ctx context.Context) (map[client.ObjectKey]Object, error) {
	if c.NodeObjectsErr != nil {
		return nil, c.NodeObjectsErr
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.nodes, nil
}

// ListNodes returns a list of StorageOS node objects.
func (c *MockClient) ListNodes(ctx context.Context) ([]client.Object, error) {
	if c.ListNodesErr != nil {
		return nil, c.ListNodesErr
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	ret := []client.Object{}
	for _, node := range c.nodes {
		health := stosv1.NodeHealthOnline
		if !node.IsHealthy() {
			health = stosv1.NodeHealthOffline
		}
		ret = append(ret, &stosv1.Node{
			ObjectMeta: metav1.ObjectMeta{Name: node.GetName()},
			Status: stosv1.NodeStatus{
				Health: health,
			},
		})
	}
	return ret, nil
}

// AddNode adds a node to the StorageOS cluster.
func (c *MockClient) AddNode(obj Object) error {
	c.mu.Lock()
	c.nodes[ObjectKeyFromObject(obj)] = obj
	c.mu.Unlock()
	return nil
}

// NodeExists returns true if the node exists.
func (c *MockClient) NodeExists(key client.ObjectKey) bool {
	c.mu.RLock()
	defer c.mu.RUnlock()
	if _, ok := c.nodes[key]; ok {
		return true
	}
	return false
}

// DeleteNode removes a node from the StorageOS cluster.
func (c *MockClient) DeleteNode(ctx context.Context, key client.ObjectKey) error {
	c.DeleteNodeCallCount[key]++
	if c.DeleteNodeErr != nil {
		return c.DeleteNodeErr
	}
	if !c.NodeExists(key) {
		return ErrNodeNotFound
	}
	c.mu.Lock()
	delete(c.nodes, key)
	c.mu.Unlock()
	return nil
}

// EnsureNodeLabels applies a set of labels to the StorageOS node.
func (c *MockClient) EnsureNodeLabels(ctx context.Context, key client.ObjectKey, labels map[string]string) error {
	if c.EnsureNodeLabelsErr != nil {
		return c.EnsureNodeLabelsErr
	}

	var errors *multierror.Error
	var newLabels = make(map[string]string)

	for k, v := range labels {
		switch {
		case !IsReservedLabel(k):
			newLabels[k] = v
		case k == ReservedLabelComputeOnly:
			newLabels[k] = v
		default:
			errors = multierror.Append(errors, ErrReservedLabelUnknown)
		}
	}

	c.mu.Lock()
	c.nodeLabels = newLabels
	c.mu.Unlock()
	return errors.ErrorOrNil()
}

// GetNodeLabels retrieves the set of labels.
func (c *MockClient) GetNodeLabels(key client.ObjectKey) (map[string]string, error) {
	if c.GetNodeLabelsErr != nil {
		return nil, c.GetNodeLabelsErr
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.nodeLabels, nil
}

// AddCommunityLicence adds a community licence to the StorageOS cluster.
func (c *MockClient) AddCommunityLicence() error {
	c.mu.Lock()
	c.licence = MockLicence{ClusterCapacityBytes: 1099511627776, Kind: "personal"}
	c.mu.Unlock()
	return nil
}

// AddEnterproseLicence adds an enterprise licence to the StorageOS cluster.
func (c *MockClient) AddEnterpriseLicence(capacity string) error {
	q := resource.MustParse(capacity)
	var int64Cap int64
	var ok bool
	if int64Cap, ok = q.AsInt64(); !ok {
		return fmt.Errorf("failed to parse capacity %s", capacity)
	}
	c.mu.Lock()
	c.licence = MockLicence{ClusterCapacityBytes: uint64(int64Cap), Kind: "premium"}
	c.mu.Unlock()
	return nil
}

// GetLicence retrieves the licence of the StorageOS cluster.
func (c *MockClient) GetLicence(ctx context.Context) (LicenceObject, error) {
	if c.GetLicenceErr != nil {
		return nil, c.GetLicenceErr
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.licence, nil
}

// AddVolume adds a volume to the StorageOS cluster.
func (c *MockClient) AddVolume(obj Object) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.volumes[ObjectKeyFromObject(obj)] = obj
	return nil
}

// DelVolume deletes a volume from the StorageOS cluster.
func (c *MockClient) DelVolume(obj Object) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	delete(c.volumes, ObjectKeyFromObject(obj))
	return nil
}

// UpdateNodeHealth sets the node health.
func (c *MockClient) UpdateNodeHealth(key client.ObjectKey, healthy bool) bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	n, ok := c.nodes[key]
	if !ok {
		return false
	}
	c.nodes[key] = MockObject{
		ID:        n.GetID(),
		Name:      n.GetName(),
		Namespace: n.GetNamespace(),
		Labels:    n.GetLabels(),
		Healthy:   healthy,
	}
	return true
}

// UpdateVolumeHealth sets the volume health.
func (c *MockClient) UpdateVolumeHealth(key client.ObjectKey, healthy bool) bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	n, ok := c.volumes[key]
	if !ok {
		return false
	}
	c.volumes[key] = MockObject{
		ID:        n.GetID(),
		Name:      n.GetName(),
		Namespace: n.GetNamespace(),
		Labels:    n.GetLabels(),
		Healthy:   healthy,
	}
	return true
}

func (c *MockClient) ListVolumes(ctx context.Context) ([]Object, error) {
	if c.ListVolumeErr != nil {
		return nil, c.ListVolumeErr
	}
	c.mu.RLock()
	defer c.mu.RUnlock()

	volumes := []Object{}
	for _, vol := range c.volumes {
		volumes = append(volumes, vol)
	}
	return volumes, nil
}

// GetVolume retrieves a volume object.
func (c *MockClient) GetVolume(ctx context.Context, key client.ObjectKey) (Object, error) {
	if c.GetVolumeErr != nil {
		return nil, c.GetVolumeErr
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	obj, ok := c.volumes[key]
	if !ok {
		return nil, ErrVolumeNotFound
	}
	return obj, nil
}

// VolumeObjects returns a map of volume objects, indexed on object key.
func (c *MockClient) VolumeObjects(ctx context.Context) (map[client.ObjectKey]Object, error) {
	if c.ListNodesErr != nil {
		return nil, c.ListNodesErr
	}
	c.mu.RLock()
	defer c.mu.RUnlock()
	return c.volumes, nil
}

// EnsureVolumeLabels applies a set of labels to the StorageOS volume.
func (c *MockClient) EnsureVolumeLabels(ctx context.Context, key client.ObjectKey, labels map[string]string) error {
	if c.EnsureVolumeLabelsErr != nil {
		return c.EnsureVolumeLabelsErr
	}

	var errors *multierror.Error
	var newLabels = make(map[string]string)

	for k, v := range labels {
		switch {
		case !IsReservedLabel(k):
			newLabels[k] = v
		case k == ReservedLabelReplicas || k == ReservedLabelFailureMode:
			newLabels[k] = v
		default:
			errors = multierror.Append(errors, ErrReservedLabelUnknown)
		}
	}

	c.mu.Lock()
	defer c.mu.Unlock()
	n, ok := c.volumes[key]
	if !ok {
		return ErrVolumeNotFound
	}
	c.volumes[key] = &MockObject{
		ID:        n.GetID(),
		Name:      n.GetName(),
		Namespace: n.GetNamespace(),
		Labels:    newLabels,
	}
	return errors.ErrorOrNil()
}

// SetExternalEndpoint sets the external endpoint on a SharedVolume.  The
// endpoint should be <host|ip>:<port>.
func (c *MockClient) SetExternalEndpoint(ctx context.Context, id string, namespace string, endpoint string) error {
	if c.SetEndpointErr != nil {
		return c.SetEndpointErr
	}
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.setExternalEndpoint != nil {
		return c.setExternalEndpoint(id, endpoint)
	}
	return nil
}

// WithSetExternalEndpoint configures setExternalEndpoint callback.
func (c *MockClient) WithSetExternalEndpoint(callback func(string, string) error) {
	c.setExternalEndpoint = callback
}

// Reset the shared volume list.
func (c *MockClient) Reset() {
	c.mu.Lock()
	c.setExternalEndpoint = nil
	c.namespaces = make(map[client.ObjectKey]Object)
	c.nodes = make(map[client.ObjectKey]Object)
	c.nodeLabels = make(map[string]string)
	c.licence = nil
	c.DeleteNamespaceCallCount = make(map[client.ObjectKey]int)
	c.DeleteNodeCallCount = make(map[client.ObjectKey]int)
	c.ListNamespacesErr = nil
	c.DeleteNamespaceErr = nil
	c.ListNodesErr = nil
	c.DeleteNodeErr = nil
	c.EnsureNodeLabelsErr = nil
	c.GetNodeLabelsErr = nil
	c.SharedVolErr = nil
	c.SharedVolsErr = nil
	c.SetEndpointErr = nil
	c.licence = nil
	c.GetLicenceErr = nil
	c.mu.Unlock()
}

// RandomSharedVolume returns a randomly generated shared volume.  Always uses default
// namespace since it will always exist.
func (c *MockClient) RandomSharedVolume() *SharedVolume {
	se := fmt.Sprintf("%d.%d.%d.%d:%d", rand.Intn(100)+1, rand.Intn(253)+1, rand.Intn(253)+1, rand.Intn(253)+1, rand.Intn(65534)+1)
	vol := stosv1.Volume{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Volume",
			APIVersion: stosv1.GroupVersion.Version,
		},
		ObjectMeta: metav1.ObjectMeta{
			Name: randomString(8),
		},
		Spec: stosv1.VolumeSpec{
			ConfigLabels: map[string]string{
				ReservedLabelK8sPVName:       "pvc-" + uuid.New().String(),
				ReservedLabelK8sPVCName:      randomString(8),
				ReservedLabelK8sPVCNamespace: "default",
			},
			NFS: stosv1.NFSConfig{
				ServiceEndpoint: &se,
			},
		},
		Status: stosv1.VolumeStatus{
			VolumeID: randomString(32),
		},
	}
	return &SharedVolume{
		Volume: &vol,
	}
}

// RandomStorageOSVolume returns a randomly generated StorageOS volume.  Always uses default
// namespace since it will always exist.
func (c *MockClient) RandomStorageOSVolume() *stosapi.Volume {
	return &stosapi.Volume{
		Id:          randomString(4),
		Name:        randomString(8),
		Description: randomString(255),
		AttachedOn:  randomString(8),
		Nfs: stosapi.NfsConfig{
			Exports: &[]stosapi.NfsExportConfig{
				{
					ExportID:   0,
					Path:       randomString(16),
					PseudoPath: randomString(16),
					Acls: []stosapi.NfsAcl{
						{
							Identity: stosapi.NfsAclIdentity{
								IdentityType: randomString(4),
								Matcher:      randomString(4),
							},
							SquashConfig: stosapi.NfsAclSquashConfig{
								Uid:    rand.Int63(),
								Gid:    rand.Int63(),
								Squash: randomString(4),
							},
							AccessLevel: randomString(4),
						},
					},
				},
			},
		},
		NamespaceID: randomString(8),
		Labels: map[string]string{
			ReservedLabelK8sPVName:       "pvc-" + uuid.New().String(),
			ReservedLabelK8sPVCName:      randomString(8),
			ReservedLabelK8sPVCNamespace: "default",
		},
		FsType:         stosapi.FSTYPE_EXT2,
		AttachmentType: stosapi.ATTACHTYPE_HOST,
		Master: stosapi.MasterDeploymentInfo{
			Id:         randomString(4),
			NodeID:     randomString(8),
			Promotable: true,
			Health:     stosapi.MASTERHEALTH_ONLINE,
		},
		Replicas: &[]stosapi.ReplicaDeploymentInfo{
			{
				Id:         randomString(4),
				NodeID:     randomString(8),
				Promotable: true,
				Health:     stosapi.REPLICAHEALTH_PROVISIONED,
			},
		},
		SizeBytes: 1,
		CreatedAt: time.Now(),
		UpdatedAt: time.Now(),
		Version:   randomString(2),
	}
}

func randomString(n int) string {
	var letterRunes = []rune("abcdefghijklmnopqrstuvwxyz")

	b := make([]rune, n)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}