diff --git a/cns/service/main.go b/cns/service/main.go index d7b9a526d5..aef1178e5b 100644 --- a/cns/service/main.go +++ b/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,11 +1694,16 @@ 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}) @@ -1706,17 +1711,51 @@ func createOrUpdateNodeInfoCRD(ctx context.Context, restConfig *rest.Config, nod 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 { diff --git a/cns/service/main_test.go b/cns/service/main_test.go index 42a64df761..b824ea8e4c 100644 --- a/cns/service/main_test.go +++ b/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) + }) + } +}