Add networking mock

Signed-off-by: Heba Elayoty <hebaelayoty@gmail.com>
virtual-kubelet · Apr 11, 2023 · 27a5140 · 27a5140
1 parent c7c4acc
commit 27a5140
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 12 deletions.
diff --git a/pkg/network/aci_network.go b/pkg/network/aci_network.go
@@ -41,6 +41,12 @@ var (
 	subnetAction   = "Microsoft.Network/virtualNetworks/subnets/action"
 )
 
+type ProviderNetworkInterface interface {
+	getSubnetClient(ctx context.Context, azConfig *auth.Config) (*aznetworkv2.SubnetsClient, error)
+	getACISubnet(ctx context.Context, subnetsClient *aznetworkv2.SubnetsClient) (aznetworkv2.Subnet, error)
+	createACISubnet(ctx context.Context, subnetsClient *aznetworkv2.SubnetsClient) error
+}
+
 type ProviderNetwork struct {
 	VnetSubscriptionID string
 	VnetName           string
@@ -112,7 +118,7 @@ func (pn *ProviderNetwork) setupNetwork(ctx context.Context, azConfig *auth.Conf
 	ctx, span := trace.StartSpan(ctx, "network.setupNetwork")
 	defer span.End()
 
-	subnetsClient, err := getSubnetClient(ctx, azConfig)
+	subnetsClient, err := pn.getSubnetClient(ctx, azConfig)
 	if err != nil {
 		return err
 	}
@@ -121,18 +127,10 @@ func (pn *ProviderNetwork) setupNetwork(ctx context.Context, azConfig *auth.Conf
 	ctxWithResp := runtime.WithCaptureResponse(ctx, &rawResponse)
 
 	createSubnet := true
-	response, err := subnetsClient.Get(ctxWithResp, pn.VnetResourceGroup, pn.VnetName, pn.SubnetName, nil)
-	var respErr *azcore.ResponseError
+	currentSubnet, err := pn.getACISubnet(ctxWithResp, subnetsClient)
 	if err != nil {
-		if errors.As(err, &respErr) && !(respErr.RawResponse.StatusCode == http.StatusNotFound) {
-			return fmt.Errorf("error while looking up subnet: %v", err)
-		}
-
-		if respErr.RawResponse.StatusCode == http.StatusNotFound && pn.SubnetCIDR == "" {
-			return fmt.Errorf("subnet '%s' is not found in vnet '%s' in resource group '%s' and subscription '%s' and subnet CIDR is not specified", pn.SubnetName, pn.VnetName, pn.VnetResourceGroup, pn.VnetSubscriptionID)
-		}
+		return err
 	}
-	currentSubnet := response.Subnet
 
 	if err == nil {
 		if currentSubnet.Properties.AddressPrefix != nil {
@@ -181,7 +179,22 @@ func (pn *ProviderNetwork) setupNetwork(ctx context.Context, azConfig *auth.Conf
 	return nil
 }
 
-func getSubnetClient(ctx context.Context, azConfig *auth.Config) (*aznetworkv2.SubnetsClient, error) {
+func (pn *ProviderNetwork) getACISubnet(ctx context.Context, subnetsClient *aznetworkv2.SubnetsClient) (aznetworkv2.Subnet, error) {
+	response, err := subnetsClient.Get(ctx, pn.VnetResourceGroup, pn.VnetName, pn.SubnetName, nil)
+	var respErr *azcore.ResponseError
+	if err != nil {
+		if errors.As(err, &respErr) && !(respErr.RawResponse.StatusCode == http.StatusNotFound) {
+			return aznetworkv2.Subnet{}, fmt.Errorf("error while looking up subnet: %v", err)
+		}
+
+		if respErr.RawResponse.StatusCode == http.StatusNotFound && pn.SubnetCIDR == "" {
+			return aznetworkv2.Subnet{}, fmt.Errorf("subnet '%s' is not found in vnet '%s' in resource group '%s' and subscription '%s' and subnet CIDR is not specified", pn.SubnetName, pn.VnetName, pn.VnetResourceGroup, pn.VnetSubscriptionID)
+		}
+	}
+	return response.Subnet, nil
+}
+
+func (pn *ProviderNetwork) getSubnetClient(ctx context.Context, azConfig *auth.Config) (*aznetworkv2.SubnetsClient, error) {
 	logger := log.G(ctx).WithField("method", "getSubnetClient")
 	ctx, span := trace.StartSpan(ctx, "network.getSubnetClient")
 	defer span.End()

diff --git a/pkg/network/mock_aci_network.go b/pkg/network/mock_aci_network.go
@@ -3,3 +3,45 @@ Copyright (c) Microsoft Corporation.
 Licensed under the Apache 2.0 license.
 */
 package network
+
+import (
+	"context"
+
+	aznetworkv2 "github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v2"
+	"github.com/virtual-kubelet/azure-aci/pkg/auth"
+)
+
+type GetSubnetClientFunc func(ctx context.Context, azConfig *auth.Config) (*aznetworkv2.SubnetsClient, error)
+type GetACISubnetFunc func(ctx context.Context, subnetsClient *aznetworkv2.SubnetsClient) (aznetworkv2.Subnet, error)
+type CreateACISubnetFunc func(ctx context.Context, subnetsClient *aznetworkv2.SubnetsClient) error
+
+type MockProviderNetwork struct {
+	MockGetSubnetClient GetSubnetClientFunc
+	MockGetACISubnet    GetACISubnetFunc
+	MockCreateACISubnet CreateACISubnetFunc
+}
+
+func NewMockACINetwork() *MockProviderNetwork {
+	return &MockProviderNetwork{}
+}
+
+func (m *MockProviderNetwork) getSubnetClient(ctx context.Context, azConfig *auth.Config) (*aznetworkv2.SubnetsClient, error) {
+	if m.MockGetSubnetClient != nil {
+		return m.MockGetSubnetClient(ctx, azConfig)
+	}
+	return nil, nil
+}
+
+func (m *MockProviderNetwork) getACISubnet(ctx context.Context, subnetsClient *aznetworkv2.SubnetsClient) (aznetworkv2.Subnet, error) {
+	if m.MockGetACISubnet != nil {
+		return m.MockGetACISubnet(ctx, subnetsClient)
+	}
+	return aznetworkv2.Subnet{}, nil
+}
+
+func (m *MockProviderNetwork) createACISubnet(ctx context.Context, subnetsClient *aznetworkv2.SubnetsClient) error {
+	if m.MockCreateACISubnet != nil {
+		return m.MockCreateACISubnet(ctx, subnetsClient)
+	}
+	return nil
+}