@@ -25,6 +25,7 @@ import (
2525 "os"
2626 "reflect"
2727 "sort"
28+ "strconv"
2829 "strings"
2930 "testing"
3031
@@ -111,6 +112,74 @@ func TestExtractAllocatableResourcesFromAsg(t *testing.T) {
111112 assert .Equal (t , (& expectedEphemeralStorage ).String (), labels ["ephemeral-storage" ].String ())
112113}
113114
115+ func TestBuildNodeFromTemplate (t * testing.T ) {
116+ awsManager := & AwsManager {}
117+ asg := & asg {AwsRef : AwsRef {Name : "test-auto-scaling-group" }}
118+ c5Instance := & InstanceType {
119+ InstanceType : "c5.xlarge" ,
120+ VCPU : 4 ,
121+ MemoryMb : 8192 ,
122+ GPU : 0 ,
123+ }
124+
125+ // Node with custom resource
126+ ephemeralStorageKey := "ephemeral-storage"
127+ ephemeralStorageValue := int64 (20 )
128+ vpcIPKey := "vpc.amazonaws.com/PrivateIPv4Address"
129+ observedNode , observedErr := awsManager .buildNodeFromTemplate (asg , & asgTemplate {
130+ InstanceType : c5Instance ,
131+ Tags : []* autoscaling.TagDescription {
132+ {
133+ Key : aws .String (fmt .Sprintf ("k8s.io/cluster-autoscaler/node-template/resources/%s" , ephemeralStorageKey )),
134+ Value : aws .String (strconv .FormatInt (ephemeralStorageValue , 10 )),
135+ },
136+ },
137+ })
138+ assert .NoError (t , observedErr )
139+ esValue , esExist := observedNode .Status .Capacity [apiv1 .ResourceName (ephemeralStorageKey )]
140+ assert .True (t , esExist )
141+ assert .Equal (t , int64 (20 ), esValue .Value ())
142+ _ , ipExist := observedNode .Status .Capacity [apiv1 .ResourceName (vpcIPKey )]
143+ assert .False (t , ipExist )
144+
145+ // Nod with labels
146+ GPULabelValue := "nvidia-telsa-v100"
147+ observedNode , observedErr = awsManager .buildNodeFromTemplate (asg , & asgTemplate {
148+ InstanceType : c5Instance ,
149+ Tags : []* autoscaling.TagDescription {
150+ {
151+ Key : aws .String (fmt .Sprintf ("k8s.io/cluster-autoscaler/node-template/label/%s" , GPULabel )),
152+ Value : aws .String (GPULabelValue ),
153+ },
154+ },
155+ })
156+ assert .NoError (t , observedErr )
157+ gpuValue , gpuLabelExist := observedNode .Labels [GPULabel ]
158+ assert .True (t , gpuLabelExist )
159+ assert .Equal (t , GPULabelValue , gpuValue )
160+
161+ // Node with taints
162+ gpuTaint := apiv1.Taint {
163+ Key : "nvidia.com/gpu" ,
164+ Value : "present" ,
165+ Effect : "NoSchedule" ,
166+ }
167+ observedNode , observedErr = awsManager .buildNodeFromTemplate (asg , & asgTemplate {
168+ InstanceType : c5Instance ,
169+ Tags : []* autoscaling.TagDescription {
170+ {
171+ Key : aws .String (fmt .Sprintf ("k8s.io/cluster-autoscaler/node-template/taint/%s" , gpuTaint .Key )),
172+ Value : aws .String (fmt .Sprintf ("%s:%s" , gpuTaint .Value , gpuTaint .Effect )),
173+ },
174+ },
175+ })
176+
177+ assert .NoError (t , observedErr )
178+ observedTaints := observedNode .Spec .Taints
179+ assert .Equal (t , 1 , len (observedTaints ))
180+ assert .Equal (t , gpuTaint , observedTaints [0 ])
181+ }
182+
114183func TestExtractLabelsFromAsg (t * testing.T ) {
115184 tags := []* autoscaling.TagDescription {
116185 {
0 commit comments