From 2862ec370c16509c20eb7aaf07ba8002520dd5c1 Mon Sep 17 00:00:00 2001
From: willie-yao <williamyao@microsoft.com>
Date: Wed, 4 Jan 2023 00:07:33 +0000
Subject: [PATCH] Add test

---
 azure/scope/managedmachinepool_test.go        | 101 ++++++++++++++++++
 .../v1beta1/azuremanagedmachinepool_types.go  |   5 +-
 2 files changed, 105 insertions(+), 1 deletion(-)

diff --git a/azure/scope/managedmachinepool_test.go b/azure/scope/managedmachinepool_test.go
index afdead1a82e2..6f1c3efe2863 100644
--- a/azure/scope/managedmachinepool_test.go
+++ b/azure/scope/managedmachinepool_test.go
@@ -624,6 +624,101 @@ func TestManagedMachinePoolScope_OSDiskType(t *testing.T) {
 	}
 }
 
+func TestManagedMachinePoolScope_KubeletDiskType(t *testing.T) {
+	scheme := runtime.NewScheme()
+	_ = expv1.AddToScheme(scheme)
+	_ = infrav1exp.AddToScheme(scheme)
+
+	cases := []struct {
+		Name     string
+		Input    ManagedMachinePoolScopeParams
+		Expected azure.ResourceSpecGetter
+	}{
+		{
+			Name: "Without KubeletDiskType",
+			Input: ManagedMachinePoolScopeParams{
+				Cluster: &clusterv1.Cluster{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "cluster1",
+						Namespace: "default",
+					},
+				},
+				ControlPlane: &infrav1exp.AzureManagedControlPlane{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "cluster1",
+						Namespace: "default",
+					},
+					Spec: infrav1exp.AzureManagedControlPlaneSpec{
+						SubscriptionID: "00000000-0000-0000-0000-000000000000",
+					},
+				},
+				ManagedMachinePool: ManagedMachinePool{
+					MachinePool:      getMachinePool("pool0"),
+					InfraMachinePool: getAzureMachinePool("pool0", infrav1exp.NodePoolModeSystem),
+				},
+			},
+			Expected: &agentpools.AgentPoolSpec{
+				Name:         "pool0",
+				SKU:          "Standard_D2s_v3",
+				Replicas:     1,
+				Mode:         "System",
+				Cluster:      "cluster1",
+				VnetSubnetID: "/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups//providers/Microsoft.Network/virtualNetworks//subnets/",
+				Headers:      map[string]string{},
+			},
+		},
+		{
+			Name: "With KubeletDiskType",
+			Input: ManagedMachinePoolScopeParams{
+				Cluster: &clusterv1.Cluster{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "cluster1",
+						Namespace: "default",
+					},
+				},
+				ControlPlane: &infrav1exp.AzureManagedControlPlane{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "cluster1",
+						Namespace: "default",
+					},
+					Spec: infrav1exp.AzureManagedControlPlaneSpec{
+						SubscriptionID: "00000000-0000-0000-0000-000000000000",
+					},
+				},
+				ManagedMachinePool: ManagedMachinePool{
+					MachinePool:      getMachinePool("pool1"),
+					InfraMachinePool: getAzureMachinePoolWithKubeletDiskType("pool1", (*infrav1exp.KubeletDiskType)(to.StringPtr("Temporary"))),
+				},
+			},
+			Expected: &agentpools.AgentPoolSpec{
+				Name:            "pool1",
+				SKU:             "Standard_D2s_v3",
+				Mode:            "User",
+				Cluster:         "cluster1",
+				Replicas:        1,
+				KubeletDiskType: (*infrav1exp.KubeletDiskType)(to.StringPtr("Temporary")),
+				VnetSubnetID:    "/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups//providers/Microsoft.Network/virtualNetworks//subnets/",
+				Headers:         map[string]string{},
+			},
+		},
+	}
+
+	for _, c := range cases {
+		c := c
+		t.Run(c.Name, func(t *testing.T) {
+			g := NewWithT(t)
+			fakeClient := fake.NewClientBuilder().WithScheme(scheme).WithObjects(c.Input.MachinePool, c.Input.InfraMachinePool, c.Input.ControlPlane).Build()
+			c.Input.Client = fakeClient
+			s, err := NewManagedMachinePoolScope(context.TODO(), c.Input)
+			g.Expect(err).To(Succeed())
+			agentPool := s.AgentPoolSpec()
+			if !reflect.DeepEqual(c.Expected, agentPool) {
+				t.Errorf("Got difference between expected result and result:\n%s", cmp.Diff(c.Expected, agentPool))
+			}
+		})
+	}
+}
+
 func getAzureMachinePool(name string, mode infrav1exp.NodePoolMode) *infrav1exp.AzureManagedMachinePool {
 	return &infrav1exp.AzureManagedMachinePool{
 		ObjectMeta: metav1.ObjectMeta{
@@ -675,6 +770,12 @@ func getAzureMachinePoolWithOsDiskType(name string, osDiskType string) *infrav1e
 	return managedPool
 }
 
+func getAzureMachinePoolWithKubeletDiskType(name string, kubeletDiskType *infrav1exp.KubeletDiskType) *infrav1exp.AzureManagedMachinePool {
+	managedPool := getAzureMachinePool(name, infrav1exp.NodePoolModeUser)
+	managedPool.Spec.KubeletDiskType = kubeletDiskType
+	return managedPool
+}
+
 func getAzureMachinePoolWithLabels(name string, nodeLabels map[string]string) *infrav1exp.AzureManagedMachinePool {
 	managedPool := getAzureMachinePool(name, infrav1exp.NodePoolModeSystem)
 	managedPool.Spec.NodeLabels = nodeLabels
diff --git a/exp/api/v1beta1/azuremanagedmachinepool_types.go b/exp/api/v1beta1/azuremanagedmachinepool_types.go
index c13a909c9f41..e6e2f5973aef 100644
--- a/exp/api/v1beta1/azuremanagedmachinepool_types.go
+++ b/exp/api/v1beta1/azuremanagedmachinepool_types.go
@@ -54,10 +54,13 @@ const (
 // TopologyManagerPolicy enumerates the values for KubeletConfig.TopologyManagerPolicy.
 type TopologyManagerPolicy string
 
+// KubeletDiskType enumerates the values for the agent pool's KubeletDiskType
 type KubeletDiskType string
 
 const (
-	KubeletDiskTypeOS        KubeletDiskType = "OS"
+	// KubeletDiskTypeOS ...
+	KubeletDiskTypeOS KubeletDiskType = "OS"
+	// KubeletDiskTypeTemporary ...
 	KubeletDiskTypeTemporary KubeletDiskType = "Temporary"
 )