blob: 87e0678f17b690aa2dd58da3cff009f8f1f0cbf6 [file] [log] [blame]
// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package cloud
import (
"context"
"crypto/rand"
"errors"
"fmt"
mrand "math/rand"
"sync"
"time"
"github.com/aws/aws-sdk-go/service/ec2"
)
// FakeAWSClient provides a fake AWS Client used to test the AWS client
// functionality.
type FakeAWSClient struct {
mu sync.RWMutex
instances map[string]*Instance
instanceTypes []*InstanceType
serviceQuotas map[serviceQuotaKey]int64
}
// serviceQuotaKey should be used as the key in the serviceQuotas map.
type serviceQuotaKey struct {
code string
service string
}
// NewFakeAWSClient crates a fake AWS client.
func NewFakeAWSClient() *FakeAWSClient {
return &FakeAWSClient{
instances: make(map[string]*Instance),
instanceTypes: []*InstanceType{
&InstanceType{"ab.large", 10},
&InstanceType{"ab.xlarge", 20},
&InstanceType{"ab.small", 30},
},
serviceQuotas: map[serviceQuotaKey]int64{
serviceQuotaKey{QuotaCodeCPUOnDemand, QuotaServiceEC2}: 384,
},
}
}
// Instance returns the `Instance` record for the requested instance. The instance record will
// return records for recently terminated instances. If an instance is not found an error will
// be returned.
func (f *FakeAWSClient) Instance(ctx context.Context, instID string) (*Instance, error) {
if ctx == nil || instID == "" {
return nil, errors.New("invalid params")
}
f.mu.RLock()
defer f.mu.RUnlock()
inst, ok := f.instances[instID]
if !ok {
return nil, errors.New("instance not found")
}
return copyInstance(inst), nil
}
// Instances retrieves all EC2 instances in a region which have not been terminated or stopped.
func (f *FakeAWSClient) RunningInstances(ctx context.Context) ([]*Instance, error) {
if ctx == nil {
return nil, errors.New("invalid params")
}
f.mu.RLock()
defer f.mu.RUnlock()
instances := make([]*Instance, 0, len(f.instances))
for _, inst := range f.instances {
if inst.State != ec2.InstanceStateNameRunning && inst.State != ec2.InstanceStateNamePending {
continue
}
instances = append(instances, copyInstance(inst))
}
return instances, nil
}
// InstanceTypesARM retrieves all EC2 instance types in a region which support the ARM64 architecture.
func (f *FakeAWSClient) InstanceTypesARM(ctx context.Context) ([]*InstanceType, error) {
if ctx == nil {
return nil, errors.New("invalid params")
}
f.mu.RLock()
defer f.mu.RUnlock()
instanceTypes := make([]*InstanceType, 0, len(f.instanceTypes))
for _, it := range f.instanceTypes {
instanceTypes = append(instanceTypes, &InstanceType{it.Type, it.CPU})
}
return instanceTypes, nil
}
// Quota retrieves the requested service quota for the service.
func (f *FakeAWSClient) Quota(ctx context.Context, service, code string) (int64, error) {
if ctx == nil || service == "" || code == "" {
return 0, errors.New("invalid params")
}
f.mu.RLock()
defer f.mu.RUnlock()
v, ok := f.serviceQuotas[serviceQuotaKey{code, service}]
if !ok {
return 0, errors.New("service quota not found")
}
return v, nil
}
// CreateInstance creates an EC2 VM instance.
func (f *FakeAWSClient) CreateInstance(ctx context.Context, config *EC2VMConfiguration) (*Instance, error) {
if ctx == nil || config == nil {
return nil, errors.New("invalid params")
}
if config.ImageID == "" {
return nil, errors.New("invalid Image ID")
}
if config.Type == "" {
return nil, errors.New("invalid Type")
}
if config.Zone == "" {
return nil, errors.New("invalid Zone")
}
f.mu.Lock()
defer f.mu.Unlock()
inst := &Instance{
CPUCount: 4,
CreatedAt: time.Now(),
Description: config.Description,
ID: fmt.Sprintf("instance-%s", randHex(10)),
IPAddressExternal: randIPv4(),
IPAddressInternal: randIPv4(),
ImageID: config.ImageID,
Name: config.Name,
SSHKeyID: config.SSHKeyID,
SecurityGroups: config.SecurityGroups,
State: ec2.InstanceStateNameRunning,
Tags: make(map[string]string),
Type: config.Type,
Zone: config.Zone,
}
for k, v := range config.Tags {
inst.Tags[k] = v
}
f.instances[inst.ID] = inst
return copyInstance(inst), nil
}
// DestroyInstances terminates EC2 VM instances.
func (f *FakeAWSClient) DestroyInstances(ctx context.Context, instIDs ...string) error {
if ctx == nil || len(instIDs) == 0 {
return errors.New("invalid params")
}
f.mu.Lock()
defer f.mu.Unlock()
for _, id := range instIDs {
inst, ok := f.instances[id]
if !ok {
return errors.New("instance not found")
}
inst.State = ec2.InstanceStateNameTerminated
}
return nil
}
// WaitUntilInstanceRunning returns when an instance has transitioned into the running state.
func (f *FakeAWSClient) WaitUntilInstanceRunning(ctx context.Context, instID string) error {
if ctx == nil || instID == "" {
return errors.New("invalid params")
}
f.mu.RLock()
defer f.mu.RUnlock()
inst, ok := f.instances[instID]
if !ok {
return errors.New("instance not found")
}
if inst.State != ec2.InstanceStateNameRunning {
return errors.New("timed out waiting for instance to enter running state")
}
return nil
}
// copyInstance copies the contents of a pointer to an instance and returns a newly created
// instance with the same data as the original instance.
func copyInstance(inst *Instance) *Instance {
i := &Instance{
CPUCount: inst.CPUCount,
CreatedAt: inst.CreatedAt,
Description: inst.Description,
ID: inst.ID,
IPAddressExternal: inst.IPAddressExternal,
IPAddressInternal: inst.IPAddressInternal,
ImageID: inst.ImageID,
Name: inst.Name,
SSHKeyID: inst.SSHKeyID,
SecurityGroups: inst.SecurityGroups,
State: inst.State,
Tags: make(map[string]string),
Type: inst.Type,
Zone: inst.Zone,
}
for k, v := range inst.Tags {
i.Tags[k] = v
}
return i
}
// randHex creates a random hex string of length n.
func randHex(n int) string {
buf := make([]byte, n/2+1)
_, _ = rand.Read(buf)
return fmt.Sprintf("%x", buf)[:n]
}
// randIPv4 creates a random IPv4 address.
func randIPv4() string {
return fmt.Sprintf("%d.%d.%d.%d", mrand.Intn(255), mrand.Intn(255), mrand.Intn(255), mrand.Intn(255))
}