From c9b9a47b942e455f489a43278a70ff7bce0b6ffe Mon Sep 17 00:00:00 2001
From: justinsb <justinsb@google.com>
Date: Sun, 4 Feb 2024 15:05:55 -0500
Subject: [PATCH] cloudmock: Implement WithContext methods for ELBv2

Also switch methods that were not passing a context.
---
 cloudmock/aws/mockelbv2/listeners.go          | 22 ++++++++++---------
 .../cloudup/awstasks/network_load_balancer.go | 12 ++++++----
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/cloudmock/aws/mockelbv2/listeners.go b/cloudmock/aws/mockelbv2/listeners.go
index b8937fae8d7ef..89fbfc7303c13 100644
--- a/cloudmock/aws/mockelbv2/listeners.go
+++ b/cloudmock/aws/mockelbv2/listeners.go
@@ -21,39 +21,41 @@ import (
 	"strings"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/elbv2"
 	"k8s.io/klog/v2"
 )
 
-func (m *MockELBV2) DescribeListeners(request *elbv2.DescribeListenersInput) (*elbv2.DescribeListenersOutput, error) {
+func (m *MockELBV2) DescribeListenersPagesWithContext(ctx aws.Context, request *elbv2.DescribeListenersInput, callback func(*elbv2.DescribeListenersOutput, bool) bool, options ...request.Option) error {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	klog.Infof("DescribeListeners v2 %v", request)
+	klog.Infof("DescribeListenersPagesWithContext v2 %v", request)
 
-	resp := &elbv2.DescribeListenersOutput{
+	page := &elbv2.DescribeListenersOutput{
 		Listeners: make([]*elbv2.Listener, 0),
 	}
 	for _, l := range m.Listeners {
 		listener := l.description
 		if aws.StringValue(request.LoadBalancerArn) == aws.StringValue(listener.LoadBalancerArn) {
-			resp.Listeners = append(resp.Listeners, &listener)
+			page.Listeners = append(page.Listeners, &listener)
 		} else {
 			for _, reqARN := range request.ListenerArns {
 				if aws.StringValue(reqARN) == aws.StringValue(listener.ListenerArn) {
-					resp.Listeners = append(resp.Listeners, &listener)
+					page.Listeners = append(page.Listeners, &listener)
 				}
 			}
 		}
 	}
-	return resp, nil
+	callback(page, true)
+	return nil
 }
 
-func (m *MockELBV2) CreateListener(request *elbv2.CreateListenerInput) (*elbv2.CreateListenerOutput, error) {
+func (m *MockELBV2) CreateListenerWithContext(ctx aws.Context, request *elbv2.CreateListenerInput, opts ...request.Option) (*elbv2.CreateListenerOutput, error) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	klog.Infof("CreateListener v2 %v", request)
+	klog.Infof("CreateListenerWithContext v2 %v", request)
 
 	l := elbv2.Listener{
 		DefaultActions:  request.DefaultActions,
@@ -96,11 +98,11 @@ func (m *MockELBV2) CreateListener(request *elbv2.CreateListenerInput) (*elbv2.C
 	return &elbv2.CreateListenerOutput{Listeners: []*elbv2.Listener{&l}}, nil
 }
 
-func (m *MockELBV2) DeleteListener(request *elbv2.DeleteListenerInput) (*elbv2.DeleteListenerOutput, error) {
+func (m *MockELBV2) DeleteListenerWithContext(ctx aws.Context, request *elbv2.DeleteListenerInput, opts ...request.Option) (*elbv2.DeleteListenerOutput, error) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
 
-	klog.Infof("DeleteListener v2 %v", request)
+	klog.Infof("DeleteListenerWithContext v2 %v", request)
 
 	lARN := aws.StringValue(request.ListenerArn)
 	if _, ok := m.Listeners[lARN]; !ok {
diff --git a/upup/pkg/fi/cloudup/awstasks/network_load_balancer.go b/upup/pkg/fi/cloudup/awstasks/network_load_balancer.go
index a099fae825c76..2fc866db73af4 100644
--- a/upup/pkg/fi/cloudup/awstasks/network_load_balancer.go
+++ b/upup/pkg/fi/cloudup/awstasks/network_load_balancer.go
@@ -196,6 +196,7 @@ func (e *NetworkLoadBalancer) getHostedZoneId() *string {
 }
 
 func (e *NetworkLoadBalancer) Find(c *fi.CloudupContext) (*NetworkLoadBalancer, error) {
+	ctx := c.Context()
 	cloud := c.T.Cloud.(awsup.AWSCloud)
 
 	lb, err := cloud.FindELBV2ByNameTag(e.Tags["Name"])
@@ -260,13 +261,16 @@ func (e *NetworkLoadBalancer) Find(c *fi.CloudupContext) (*NetworkLoadBalancer,
 		request := &elbv2.DescribeListenersInput{
 			LoadBalancerArn: loadBalancerArn,
 		}
-		response, err := cloud.ELBV2().DescribeListeners(request)
-		if err != nil {
-			return nil, fmt.Errorf("error querying for NLB listeners :%v", err)
+		var listeners []*elbv2.Listener
+		if err := cloud.ELBV2().DescribeListenersPagesWithContext(ctx, request, func(page *elbv2.DescribeListenersOutput, lastPage bool) bool {
+			listeners = append(listeners, page.Listeners...)
+			return true
+		}); err != nil {
+			return nil, fmt.Errorf("listing NLB listeners: %w", err)
 		}
 
 		actual.TargetGroups = []*TargetGroup{}
-		for _, l := range response.Listeners {
+		for _, l := range listeners {
 			// This will need to be rearranged when we recognized multiple listeners and target groups per NLB
 			if len(l.DefaultActions) > 0 {
 				targetGroupARN := l.DefaultActions[0].TargetGroupArn