Merge pull request #154 from qonto/fix-crash-on-ec2-type-missing-info

Fix crash on ec2 type missing info
qonto · Mar 26, 2024 · 9e8fcb8 · 9e8fcb8
2 parents be0e433 + 9087b58
commit 9e8fcb8
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 14 deletions.
diff --git a/internal/app/ec2/ec2.go b/internal/app/ec2/ec2.go
@@ -89,13 +89,23 @@ func (e *EC2Fetcher) GetDBInstanceTypeInformation(instanceTypes []string) (Metri
 		e.statistics.EC2ApiCall++
 
 		for _, i := range resp.InstanceTypes {
-			instanceName := addDBPrefix(string(i.InstanceType))
-			metrics[instanceName] = EC2InstanceMetrics{
-				Vcpu:              aws.ToInt32(i.VCpuInfo.DefaultVCpus),
-				MaximumIops:       aws.ToInt32(i.EbsInfo.EbsOptimizedInfo.MaximumIops),
-				MaximumThroughput: converter.MegaBytesToBytes(aws.ToFloat64(i.EbsInfo.EbsOptimizedInfo.MaximumThroughputInMBps)),
-				Memory:            converter.MegaBytesToBytes(aws.ToInt64(i.MemoryInfo.SizeInMiB)),
+			instanceMetrics := EC2InstanceMetrics{}
+
+			if i.VCpuInfo != nil {
+				instanceMetrics.Vcpu = aws.ToInt32(i.VCpuInfo.DefaultVCpus)
+			}
+
+			if i.MemoryInfo != nil {
+				instanceMetrics.Memory = converter.MegaBytesToBytes(aws.ToInt64(i.MemoryInfo.SizeInMiB))
 			}
+
+			if i.EbsInfo != nil && i.EbsInfo.EbsOptimizedInfo != nil {
+				instanceMetrics.MaximumIops = aws.ToInt32(i.EbsInfo.EbsOptimizedInfo.MaximumIops)
+				instanceMetrics.MaximumThroughput = converter.MegaBytesToBytes(aws.ToFloat64(i.EbsInfo.EbsOptimizedInfo.MaximumThroughputInMBps))
+			}
+
+			instanceName := addDBPrefix(string(i.InstanceType))
+			metrics[instanceName] = instanceMetrics
 		}
 
 		instanceTypeSpan.SetStatus(codes.Ok, "metrics fetched")

diff --git a/internal/app/ec2/ec2_test.go b/internal/app/ec2/ec2_test.go
@@ -15,18 +15,57 @@ func TestGetDBInstanceTypeInformation(t *testing.T) {
 	context := context.TODO()
 	client := mock.EC2Client{}
 
-	instanceTypes := []string{"db.t3.large", "db.t3.small"}
+	testCases := []struct {
+		instanceType      string
+		vCPU              int32
+		memory            int64
+		maximumIops       int32
+		maximumThroughput float64
+	}{
+		{
+			instanceType:      "t3.large",
+			vCPU:              mock.InstanceT3Large.Vcpu,
+			memory:            converter.MegaBytesToBytes(mock.InstanceT3Large.Memory),
+			maximumIops:       mock.InstanceT3Large.MaximumIops,
+			maximumThroughput: converter.MegaBytesToBytes(mock.InstanceT3Large.MaximumThroughput),
+		},
+		{
+			instanceType:      "t3.small",
+			vCPU:              mock.InstanceT3Small.Vcpu,
+			memory:            converter.MegaBytesToBytes(mock.InstanceT3Small.Memory),
+			maximumIops:       mock.InstanceT3Small.MaximumIops,
+			maximumThroughput: converter.MegaBytesToBytes(mock.InstanceT3Small.MaximumThroughput),
+		},
+		{
+			instanceType:      "t2.small",
+			vCPU:              mock.InstanceT2Small.Vcpu,
+			memory:            converter.MegaBytesToBytes(mock.InstanceT2Small.Memory),
+			maximumIops:       0, // Don't have Maximum IOPS for non EBS optimized instances
+			maximumThroughput: 0, // Don't have Maximum throughput for non EBS optimized instances
+		},
+	}
+	expectedAPICalls := float64(1)
+
+	instanceTypes := make([]string, len(testCases))
+	for i, instance := range testCases {
+		instanceTypes[i] = instance.instanceType
+	}
+
 	fetcher := ec2.NewFetcher(context, client)
 	result, err := fetcher.GetDBInstanceTypeInformation(instanceTypes)
 
 	require.NoError(t, err, "GetDBInstanceTypeInformation must succeed")
-	assert.Equal(t, mock.InstanceT3Large.Vcpu, result.Instances["db.t3.large"].Vcpu, "vCPU don't match")
-	assert.Equal(t, converter.MegaBytesToBytes(mock.InstanceT3Large.Memory), result.Instances["db.t3.large"].Memory, "Memory don't match")
-	assert.Equal(t, mock.InstanceT3Large.MaximumIops, result.Instances["db.t3.large"].MaximumIops, "MaximumThroughput don't match")
-	assert.Equal(t, converter.MegaBytesToBytes(mock.InstanceT3Large.MaximumThroughput), result.Instances["db.t3.large"].MaximumThroughput, "MaximumThroughput don't match")
+	assert.Equal(t, expectedAPICalls, fetcher.GetStatistics().EC2ApiCall, "EC2 API calls don't match")
 
-	assert.Equal(t, mock.InstanceT3Small.Vcpu, result.Instances["db.t3.small"].Vcpu, "vCPU don't match")
-	assert.Equal(t, converter.MegaBytesToBytes(mock.InstanceT3Small.Memory), result.Instances["db.t3.small"].Memory, "Memory don't match")
+	for _, tc := range testCases {
+		testName := "Test " + tc.instanceType
+		t.Run(testName, func(t *testing.T) {
+			instance := result.Instances["db."+tc.instanceType]
 
-	assert.Equal(t, float64(1), fetcher.GetStatistics().EC2ApiCall, "EC2 API call don't match")
+			assert.Equal(t, tc.vCPU, instance.Vcpu, "vCPU don't match")
+			assert.Equal(t, tc.memory, instance.Memory, "Memory don't match")
+			assert.Equal(t, tc.maximumIops, instance.MaximumIops, "Maximum IOPS don't match")
+			assert.Equal(t, tc.maximumThroughput, instance.MaximumThroughput, "Maximum throughput don't match")
+		})
+	}
 }
diff --git a/internal/app/ec2/mock/ec2.go b/internal/app/ec2/mock/ec2.go
@@ -25,6 +25,12 @@ var InstanceT3Small = ec2.EC2InstanceMetrics{
 	Vcpu:              2,
 }
 
+//nolint:golint,gomnd
+var InstanceT2Small = ec2.EC2InstanceMetrics{
+	Memory: 2,
+	Vcpu:   1,
+}
+
 type EC2Client struct{}
 
 func (m EC2Client) DescribeInstanceTypes(ctx context.Context, input *aws_ec2.DescribeInstanceTypesInput, optFns ...func(*aws_ec2.Options)) (*aws_ec2.DescribeInstanceTypesOutput, error) {
@@ -53,6 +59,12 @@ func (m EC2Client) DescribeInstanceTypes(ctx context.Context, input *aws_ec2.Des
 					MaximumThroughputInMBps: &InstanceT3Small.MaximumThroughput,
 				}},
 			})
+		case "t2.small":
+			instances = append(instances, aws_ec2_types.InstanceTypeInfo{
+				InstanceType: instanceType,
+				VCpuInfo:     &aws_ec2_types.VCpuInfo{DefaultVCpus: &InstanceT2Small.Vcpu},
+				MemoryInfo:   &aws_ec2_types.MemoryInfo{SizeInMiB: &InstanceT2Small.Memory},
+			})
 		}
 	}