Skip to content

Commit

Permalink
cloudmock: Add context functions to mock
Browse files Browse the repository at this point in the history
  • Loading branch information
justinsb committed Oct 25, 2023
1 parent c4045fb commit d4d9e8a
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 34 deletions.
27 changes: 11 additions & 16 deletions cloudmock/aws/mockautoscaling/group.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ limitations under the License.
package mockautoscaling

import (
"context"
"fmt"
"time"

Expand Down Expand Up @@ -44,7 +45,7 @@ func (m *MockAutoscaling) AttachInstances(input *autoscaling.AttachInstancesInpu
return &autoscaling.AttachInstancesOutput{}, nil
}

func (m *MockAutoscaling) CreateAutoScalingGroup(input *autoscaling.CreateAutoScalingGroupInput) (*autoscaling.CreateAutoScalingGroupOutput, error) {
func (m *MockAutoscaling) CreateAutoScalingGroupWithContext(ctx aws.Context, input *autoscaling.CreateAutoScalingGroupInput, options ...request.Option) (*autoscaling.CreateAutoScalingGroupOutput, error) {
m.mutex.Lock()
defer m.mutex.Unlock()

Expand Down Expand Up @@ -102,7 +103,7 @@ func (m *MockAutoscaling) CreateAutoScalingGroup(input *autoscaling.CreateAutoSc
return &autoscaling.CreateAutoScalingGroupOutput{}, nil
}

func (m *MockAutoscaling) UpdateAutoScalingGroup(request *autoscaling.UpdateAutoScalingGroupInput) (*autoscaling.UpdateAutoScalingGroupOutput, error) {
func (m *MockAutoscaling) UpdateAutoScalingGroupWithContext(ctx context.Context, request *autoscaling.UpdateAutoScalingGroupInput, opts ...request.Option) (*autoscaling.UpdateAutoScalingGroupOutput, error) {
m.mutex.Lock()
defer m.mutex.Unlock()
klog.V(2).Infof("Mock UpdateAutoScalingGroup %v", request)
Expand Down Expand Up @@ -163,7 +164,7 @@ func (m *MockAutoscaling) UpdateAutoScalingGroup(request *autoscaling.UpdateAuto
return &autoscaling.UpdateAutoScalingGroupOutput{}, nil
}

func (m *MockAutoscaling) EnableMetricsCollection(request *autoscaling.EnableMetricsCollectionInput) (*autoscaling.EnableMetricsCollectionOutput, error) {
func (m *MockAutoscaling) EnableMetricsCollectionWithContext(ctx aws.Context, request *autoscaling.EnableMetricsCollectionInput, opts ...request.Option) (*autoscaling.EnableMetricsCollectionOutput, error) {
m.mutex.Lock()
defer m.mutex.Unlock()

Expand Down Expand Up @@ -195,7 +196,7 @@ func (m *MockAutoscaling) EnableMetricsCollection(request *autoscaling.EnableMet
return response, nil
}

func (m *MockAutoscaling) SuspendProcesses(input *autoscaling.ScalingProcessQuery) (*autoscaling.SuspendProcessesOutput, error) {
func (m *MockAutoscaling) SuspendProcessesWithContext(ctx aws.Context, input *autoscaling.ScalingProcessQuery, opts ...request.Option) (*autoscaling.SuspendProcessesOutput, error) {
m.mutex.Lock()
defer m.mutex.Unlock()

Expand Down Expand Up @@ -290,7 +291,7 @@ func (m *MockAutoscaling) DescribeAutoScalingGroupsRequest(*autoscaling.Describe
return nil, nil
}

func (m *MockAutoscaling) DescribeAutoScalingGroupsPages(request *autoscaling.DescribeAutoScalingGroupsInput, callback func(*autoscaling.DescribeAutoScalingGroupsOutput, bool) bool) error {
func (m *MockAutoscaling) DescribeAutoScalingGroupsPagesWithContext(ctx aws.Context, request *autoscaling.DescribeAutoScalingGroupsInput, callback func(*autoscaling.DescribeAutoScalingGroupsOutput, bool) bool, options ...request.Option) error {
if request.MaxRecords != nil {
klog.Fatalf("MaxRecords not implemented")
}
Expand All @@ -309,12 +310,11 @@ func (m *MockAutoscaling) DescribeAutoScalingGroupsPages(request *autoscaling.De
return nil
}

func (m *MockAutoscaling) DescribeAutoScalingGroupsPagesWithContext(aws.Context, *autoscaling.DescribeAutoScalingGroupsInput, func(*autoscaling.DescribeAutoScalingGroupsOutput, bool) bool, ...request.Option) error {
klog.Fatalf("Not implemented")
return nil
func (m *MockAutoscaling) DescribeAutoScalingGroupsPages(request *autoscaling.DescribeAutoScalingGroupsInput, callback func(*autoscaling.DescribeAutoScalingGroupsOutput, bool) bool) error {
return m.DescribeAutoScalingGroupsPagesWithContext(context.TODO(), request, callback)
}

func (m *MockAutoscaling) DeleteAutoScalingGroup(request *autoscaling.DeleteAutoScalingGroupInput) (*autoscaling.DeleteAutoScalingGroupOutput, error) {
func (m *MockAutoscaling) DeleteAutoScalingGroupWithContext(ctx aws.Context, request *autoscaling.DeleteAutoScalingGroupInput, options ...request.Option) (*autoscaling.DeleteAutoScalingGroupOutput, error) {
m.mutex.Lock()
defer m.mutex.Unlock()

Expand All @@ -330,17 +330,12 @@ func (m *MockAutoscaling) DeleteAutoScalingGroup(request *autoscaling.DeleteAuto
return &autoscaling.DeleteAutoScalingGroupOutput{}, nil
}

func (m *MockAutoscaling) DeleteAutoScalingGroupWithContext(aws.Context, *autoscaling.DeleteAutoScalingGroupInput, ...request.Option) (*autoscaling.DeleteAutoScalingGroupOutput, error) {
klog.Fatalf("Not implemented")
return nil, nil
}

func (m *MockAutoscaling) DeleteAutoScalingGroupRequest(*autoscaling.DeleteAutoScalingGroupInput) (*request.Request, *autoscaling.DeleteAutoScalingGroupOutput) {
klog.Fatalf("Not implemented")
return nil, nil
}

func (m *MockAutoscaling) PutLifecycleHook(input *autoscaling.PutLifecycleHookInput) (*autoscaling.PutLifecycleHookOutput, error) {
func (m *MockAutoscaling) PutLifecycleHookWithContext(ctx aws.Context, input *autoscaling.PutLifecycleHookInput, options ...request.Option) (*autoscaling.PutLifecycleHookOutput, error) {
m.mutex.Lock()
defer m.mutex.Unlock()
hook := &autoscaling.LifecycleHook{
Expand All @@ -364,7 +359,7 @@ func (m *MockAutoscaling) PutLifecycleHook(input *autoscaling.PutLifecycleHookIn
return &autoscaling.PutLifecycleHookOutput{}, nil
}

func (m *MockAutoscaling) DescribeLifecycleHooks(input *autoscaling.DescribeLifecycleHooksInput) (*autoscaling.DescribeLifecycleHooksOutput, error) {
func (m *MockAutoscaling) DescribeLifecycleHooksWithContext(ctx aws.Context, input *autoscaling.DescribeLifecycleHooksInput, options ...request.Option) (*autoscaling.DescribeLifecycleHooksOutput, error) {
m.mutex.Lock()
defer m.mutex.Unlock()

Expand Down
7 changes: 1 addition & 6 deletions cloudmock/aws/mockautoscaling/tags.go
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ func (m *MockAutoscaling) DescribeTagsRequest(*autoscaling.DescribeTagsInput) (*
return nil, nil
}

func (m *MockAutoscaling) DescribeTagsPages(request *autoscaling.DescribeTagsInput, callback func(*autoscaling.DescribeTagsOutput, bool) bool) error {
func (m *MockAutoscaling) DescribeTagsPagesWithContext(ctx aws.Context, request *autoscaling.DescribeTagsInput, callback func(*autoscaling.DescribeTagsOutput, bool) bool, options ...request.Option) error {
// For the mock, we just send everything in one page
page, err := m.DescribeTags(request)
if err != nil {
Expand All @@ -83,8 +83,3 @@ func (m *MockAutoscaling) DescribeTagsPages(request *autoscaling.DescribeTagsInp

return nil
}

func (m *MockAutoscaling) DescribeTagsPagesWithContext(aws.Context, *autoscaling.DescribeTagsInput, func(*autoscaling.DescribeTagsOutput, bool) bool, ...request.Option) error {
klog.Fatalf("Not implemented")
return nil
}
8 changes: 6 additions & 2 deletions cloudmock/aws/mockautoscaling/warmpool.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,13 @@ limitations under the License.

package mockautoscaling

import "github.com/aws/aws-sdk-go/service/autoscaling"
import (
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/request"
"github.com/aws/aws-sdk-go/service/autoscaling"
)

func (m *MockAutoscaling) DescribeWarmPool(input *autoscaling.DescribeWarmPoolInput) (*autoscaling.DescribeWarmPoolOutput, error) {
func (m *MockAutoscaling) DescribeWarmPoolWithContext(ctx aws.Context, input *autoscaling.DescribeWarmPoolInput, options ...request.Option) (*autoscaling.DescribeWarmPoolOutput, error) {
instances, found := m.WarmPoolInstances[*input.AutoScalingGroupName]
if !found {
return &autoscaling.DescribeWarmPoolOutput{}, nil
Expand Down
18 changes: 8 additions & 10 deletions pkg/instancegroups/rollingupdate_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,8 @@ func (v *assertNotCalledClusterValidator) Validate() (*validation.ValidationClus
}

func makeGroup(groups map[string]*cloudinstances.CloudInstanceGroup, k8sClient kubernetes.Interface, cloud awsup.AWSCloud, name string, role kopsapi.InstanceGroupRole, count int, needUpdate int) {
ctx := context.TODO()

fakeClient := k8sClient.(*fake.Clientset)

group := &cloudinstances.CloudInstanceGroup{
Expand All @@ -154,7 +156,7 @@ func makeGroup(groups map[string]*cloudinstances.CloudInstanceGroup, k8sClient k
}
groups[name] = group

cloud.Autoscaling().CreateAutoScalingGroup(&autoscaling.CreateAutoScalingGroupInput{
cloud.Autoscaling().CreateAutoScalingGroupWithContext(ctx, &autoscaling.CreateAutoScalingGroupInput{
AutoScalingGroupName: aws.String(name),
DesiredCapacity: aws.Int64(int64(count)),
MinSize: aws.Int64(1),
Expand Down Expand Up @@ -939,7 +941,7 @@ type disabledSurgeTest struct {
numDetached int
}

func (m *disabledSurgeTest) DetachInstances(input *autoscaling.DetachInstancesInput) (*autoscaling.DetachInstancesOutput, error) {
func (m *disabledSurgeTest) DetachInstancesWithContext(ctx context.Context, input *autoscaling.DetachInstancesInput, option ...request.Option) (*autoscaling.DetachInstancesOutput, error) {
m.mutex.Lock()
defer m.mutex.Unlock()

Expand All @@ -950,10 +952,6 @@ func (m *disabledSurgeTest) DetachInstances(input *autoscaling.DetachInstancesIn
return &autoscaling.DetachInstancesOutput{}, nil
}

func (m *disabledSurgeTest) DetachInstancesWithContext(ctx context.Context, input *autoscaling.DetachInstancesInput, option ...request.Option) (*autoscaling.DetachInstancesOutput, error) {
return m.DetachInstances(input)
}

func TestRollingUpdateDisabledSurge(t *testing.T) {
c, cloud := getTestSetup()

Expand Down Expand Up @@ -1238,7 +1236,7 @@ type concurrentTestAutoscaling struct {
ConcurrentTest *concurrentTest
}

func (m *concurrentTestAutoscaling) DetachInstances(input *autoscaling.DetachInstancesInput) (*autoscaling.DetachInstancesOutput, error) {
func (m *concurrentTestAutoscaling) DetachInstancesWithContext(ctx context.Context, input *autoscaling.DetachInstancesInput, option ...request.Option) (*autoscaling.DetachInstancesOutput, error) {
m.ConcurrentTest.mutex.Lock()
defer m.ConcurrentTest.mutex.Unlock()

Expand Down Expand Up @@ -1320,7 +1318,7 @@ type countDetach struct {
Count int
}

func (c *countDetach) DetachInstances(input *autoscaling.DetachInstancesInput) (*autoscaling.DetachInstancesOutput, error) {
func (c *countDetach) DetachInstancesWithContext(ctx context.Context, input *autoscaling.DetachInstancesInput, option ...request.Option) (*autoscaling.DetachInstancesOutput, error) {
c.Count += len(input.InstanceIds)
return &autoscaling.DetachInstancesOutput{}, nil
}
Expand Down Expand Up @@ -1350,7 +1348,7 @@ type failDetachAutoscaling struct {
autoscalingiface.AutoScalingAPI
}

func (m *failDetachAutoscaling) DetachInstances(input *autoscaling.DetachInstancesInput) (*autoscaling.DetachInstancesOutput, error) {
func (m *failDetachAutoscaling) DetachInstancesWithContext(ctx context.Context, input *autoscaling.DetachInstancesInput, option ...request.Option) (*autoscaling.DetachInstancesOutput, error) {
return nil, fmt.Errorf("testing error")
}

Expand Down Expand Up @@ -1458,7 +1456,7 @@ type alreadyDetachedTestAutoscaling struct {
AlreadyDetachedTest *alreadyDetachedTest
}

func (m *alreadyDetachedTestAutoscaling) DetachInstances(input *autoscaling.DetachInstancesInput) (*autoscaling.DetachInstancesOutput, error) {
func (m *alreadyDetachedTestAutoscaling) DetachInstancesWithContext(ctx aws.Context, input *autoscaling.DetachInstancesInput, options ...request.Option) (*autoscaling.DetachInstancesOutput, error) {
m.AlreadyDetachedTest.mutex.Lock()
defer m.AlreadyDetachedTest.mutex.Unlock()

Expand Down

0 comments on commit d4d9e8a

Please sign in to comment.