Add HomeAZ info to NodeInfo CRD upon creation

Author: timraymond

cns/service/main.go

@@ -1419,7 +1419,7 @@ func InitializeCRDState(ctx context.Context, z *zap.Logger, httpRestService cns.
 	if _, ok := node.Labels[configuration.LabelNodeSwiftV2]; ok {
 		cnsconfig.EnableSwiftV2 = true
 		cnsconfig.WatchPods = true
-		if nodeInfoErr := createOrUpdateNodeInfoCRD(ctx, kubeConfig, node); nodeInfoErr != nil {
+		if nodeInfoErr := createOrUpdateNodeInfoCRD(ctx, z, kubeConfig, node); nodeInfoErr != nil {
 			return errors.Wrap(nodeInfoErr, "error creating or updating nodeinfo crd")
 		}
 	}
@@ -1694,29 +1694,68 @@ func getPodInfoByIPProvider(
 
 // createOrUpdateNodeInfoCRD polls imds to learn the VM Unique ID and then creates or updates the NodeInfo CRD
 // with that vm unique ID
-func createOrUpdateNodeInfoCRD(ctx context.Context, restConfig *rest.Config, node *corev1.Node) error {
+func createOrUpdateNodeInfoCRD(ctx context.Context, z *zap.Logger, restConfig *rest.Config, node *corev1.Node) error {
 	imdsCli := imds.NewClient()
-	vmUniqueID, err := imdsCli.GetVMUniqueID(ctx)
+
+	nmaConfig, err := nmagent.NewConfig("")
 	if err != nil {
-		return errors.Wrap(err, "error getting vm unique ID from imds")
+		return errors.Wrap(err, "failed to create nmagent config")
+	}
+	nmaCli, err := nmagent.NewClient(nmaConfig)
+	if err != nil {
+		return errors.Wrap(err, "failed to create nmagent client")
 	}
 
 	directcli, err := client.New(restConfig, client.Options{Scheme: multitenancy.Scheme})
 	if err != nil {
 		return errors.Wrap(err, "failed to create ctrl client")
 	}
 
-	nodeInfoCli := multitenancy.NodeInfoClient{
+	nodeInfoCli := &multitenancy.NodeInfoClient{
 		Cli: directcli,
 	}
 
+	return buildAndCreateNodeInfo(ctx, z, imdsCli, nmaCli, nodeInfoCli, node)
+}
+
+// VMUniqueIDGetter is an interface for getting VM unique ID from IMDS
+type VMUniqueIDGetter interface {
+	GetVMUniqueID(ctx context.Context) (string, error)
+}
+
+// HomeAzGetter is an interface for getting HomeAZ from NMAgent
+type HomeAzGetter interface {
+	GetHomeAz(ctx context.Context) (nmagent.AzResponse, error)
+}
+
+// NodeInfoCreator is an interface for creating/updating NodeInfo CRD
+type NodeInfoCreator interface {
+	CreateOrUpdate(ctx context.Context, nodeInfo *mtv1alpha1.NodeInfo, fieldOwner string) error
+}
+
+// buildAndCreateNodeInfo builds the NodeInfo spec with VMUniqueID and HomeAZ and creates/updates the CRD
+func buildAndCreateNodeInfo(ctx context.Context, z *zap.Logger, imdsCli VMUniqueIDGetter, nmaCli HomeAzGetter, nodeInfoCli NodeInfoCreator, node *corev1.Node) error {
+	vmUniqueID, err := imdsCli.GetVMUniqueID(ctx)
+	if err != nil {
+		return errors.Wrap(err, "error getting vm unique ID from imds")
+	}
+
+	nodeInfoSpec := mtv1alpha1.NodeInfoSpec{
+		VMUniqueID: vmUniqueID,
+	}
+
+	homeAzResponse, err := nmaCli.GetHomeAz(ctx)
+	if err != nil {
+		z.Warn("failed to get HomeAZ from nmagent", zap.Error(err))
+	} else if homeAzResponse.HomeAz > 0 {
+		nodeInfoSpec.HomeAZ = fmt.Sprintf("AZ%02d", homeAzResponse.HomeAz)
+	}
+
 	nodeInfo := &mtv1alpha1.NodeInfo{
 		ObjectMeta: metav1.ObjectMeta{
 			Name: node.Name,
 		},
-		Spec: mtv1alpha1.NodeInfoSpec{
-			VMUniqueID: vmUniqueID,
-		},
+		Spec: nodeInfoSpec,
 	}
 
 	if err := controllerutil.SetOwnerReference(node, nodeInfo, multitenancy.Scheme); err != nil {

cns/service/main_test.go

@@ -10,7 +10,14 @@ import (
 	"github.com/Azure/azure-container-networking/cns"
 	"github.com/Azure/azure-container-networking/cns/fakes"
 	"github.com/Azure/azure-container-networking/cns/logger"
+	mtv1alpha1 "github.com/Azure/azure-container-networking/crd/multitenancy/api/v1alpha1"
+	"github.com/Azure/azure-container-networking/nmagent"
+	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 // MockHTTPClient is a mock implementation of HTTPClient
@@ -69,3 +76,149 @@ func TestSendRegisterNodeRequest_StatusAccepted(t *testing.T) {
 
 	assert.Error(t, sendRegisterNodeRequest(ctx, mockClient, httpServiceFake, nodeRegisterReq, url))
 }
+
+// mockIMDSClient is a mock implementation of the VMUniqueIDGetter interface
+type mockIMDSClient struct {
+	vmUniqueID string
+	err        error
+}
+
+func (m *mockIMDSClient) GetVMUniqueID(_ context.Context) (string, error) {
+	return m.vmUniqueID, m.err
+}
+
+// mockNMAgentClient is a mock implementation of the HomeAzGetter interface
+type mockNMAgentClient struct {
+	homeAzResponse nmagent.AzResponse
+	err            error
+}
+
+func (m *mockNMAgentClient) GetHomeAz(_ context.Context) (nmagent.AzResponse, error) {
+	return m.homeAzResponse, m.err
+}
+
+// mockNodeInfoClient is a mock implementation of the NodeInfoClient interface
+type mockNodeInfoClient struct {
+	createdNodeInfo *mtv1alpha1.NodeInfo
+	err             error
+}
+
+func (m *mockNodeInfoClient) CreateOrUpdate(_ context.Context, nodeInfo *mtv1alpha1.NodeInfo, _ string) error {
+	m.createdNodeInfo = nodeInfo
+	return m.err
+}
+
+func TestBuildNodeInfoSpec_WithHomeAZ(t *testing.T) {
+	tests := []struct {
+		name            string
+		vmUniqueID      string
+		vmUniqueIDErr   error
+		homeAzResponse  nmagent.AzResponse
+		homeAzErr       error
+		expectedSpec    mtv1alpha1.NodeInfoSpec
+		expectedNodeErr bool
+	}{
+		{
+			name:           "success with HomeAZ zone 1",
+			vmUniqueID:     "test-vm-unique-id",
+			vmUniqueIDErr:  nil,
+			homeAzResponse: nmagent.AzResponse{HomeAz: 1},
+			homeAzErr:      nil,
+			expectedSpec: mtv1alpha1.NodeInfoSpec{
+				VMUniqueID: "test-vm-unique-id",
+				HomeAZ:     "AZ01",
+			},
+			expectedNodeErr: false,
+		},
+		{
+			name:           "success with HomeAZ zone 2",
+			vmUniqueID:     "another-vm-id",
+			vmUniqueIDErr:  nil,
+			homeAzResponse: nmagent.AzResponse{HomeAz: 2},
+			homeAzErr:      nil,
+			expectedSpec: mtv1alpha1.NodeInfoSpec{
+				VMUniqueID: "another-vm-id",
+				HomeAZ:     "AZ02",
+			},
+			expectedNodeErr: false,
+		},
+		{
+			name:           "success with HomeAZ zone 10",
+			vmUniqueID:     "vm-id-zone10",
+			vmUniqueIDErr:  nil,
+			homeAzResponse: nmagent.AzResponse{HomeAz: 10},
+			homeAzErr:      nil,
+			expectedSpec: mtv1alpha1.NodeInfoSpec{
+				VMUniqueID: "vm-id-zone10",
+				HomeAZ:     "AZ10",
+			},
+			expectedNodeErr: false,
+		},
+		{
+			name:           "HomeAZ not available - should succeed without HomeAZ",
+			vmUniqueID:     "test-vm-id",
+			vmUniqueIDErr:  nil,
+			homeAzResponse: nmagent.AzResponse{},
+			homeAzErr:      errors.New("nmagent HomeAZ not available"),
+			expectedSpec: mtv1alpha1.NodeInfoSpec{
+				VMUniqueID: "test-vm-id",
+				HomeAZ:     "", // HomeAZ should be empty when not available
+			},
+			expectedNodeErr: false,
+		},
+		{
+			name:            "IMDS error - should fail",
+			vmUniqueID:      "",
+			vmUniqueIDErr:   errors.New("imds error"),
+			homeAzResponse:  nmagent.AzResponse{HomeAz: 1},
+			homeAzErr:       nil,
+			expectedSpec:    mtv1alpha1.NodeInfoSpec{},
+			expectedNodeErr: true,
+		},
+		{
+			name:           "HomeAZ zone 0 - should be treated as not available",
+			vmUniqueID:     "test-vm-id",
+			vmUniqueIDErr:  nil,
+			homeAzResponse: nmagent.AzResponse{HomeAz: 0},
+			homeAzErr:      nil,
+			expectedSpec: mtv1alpha1.NodeInfoSpec{
+				VMUniqueID: "test-vm-id",
+				HomeAZ:     "",
+			},
+			expectedNodeErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			z := zap.NewNop()
+
+			imdsCli := &mockIMDSClient{
+				vmUniqueID: tt.vmUniqueID,
+				err:        tt.vmUniqueIDErr,
+			}
+			nmaCli := &mockNMAgentClient{
+				homeAzResponse: tt.homeAzResponse,
+				err:            tt.homeAzErr,
+			}
+			nodeInfoCli := &mockNodeInfoClient{}
+			node := &corev1.Node{
+				ObjectMeta: metav1.ObjectMeta{
+					Name: "test-node",
+					UID:  "test-uid",
+				},
+			}
+
+			err := buildAndCreateNodeInfo(context.Background(), z, imdsCli, nmaCli, nodeInfoCli, node)
+
+			if tt.expectedNodeErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			require.NotNil(t, nodeInfoCli.createdNodeInfo)
+			assert.Equal(t, tt.expectedSpec.VMUniqueID, nodeInfoCli.createdNodeInfo.Spec.VMUniqueID)
+			assert.Equal(t, tt.expectedSpec.HomeAZ, nodeInfoCli.createdNodeInfo.Spec.HomeAZ)
+		})
+	}
+}