Merge pull request tailwarden#1090 from Azanul/util-funcs

Util funcs
kavibritto · Oct 15, 2023 · 9a6974f · 9a6974f
2 parents 6215b4c + a5a51a6
commit 9a6974f
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 10 deletions.
diff --git a/providers/aws/lambda/functions.go b/providers/aws/lambda/functions.go
@@ -48,7 +48,7 @@ func Functions(ctx context.Context, client providers.ProviderClient) ([]models.R
 		return resources, err
 	}
 
-	priceMap, err := awsUtils.GetPriceMap(pricingOutput)
+	priceMap, err := awsUtils.GetPriceMap(pricingOutput, "group")
 	if err != nil {
 		log.Errorf("ERROR: Failed to calculate cost per month: %v", err)
 		return resources, err

diff --git a/providers/aws/s3/buckets.go b/providers/aws/s3/buckets.go
@@ -45,7 +45,7 @@ func Buckets(ctx context.Context, client ProviderClient) ([]Resource, error) {
 		return resources, err
 	}
 
-	priceMap, err := awsUtils.GetPriceMap(pricingOutput)
+	priceMap, err := awsUtils.GetPriceMap(pricingOutput, "group")
 	if err != nil {
 		log.Errorf("ERROR: Failed to calculate cost per month: %v", err)
 		return resources, err

diff --git a/providers/aws/utils/utils.go b/providers/aws/utils/utils.go
@@ -11,8 +11,12 @@ import (
 type ProductEntry struct {
 	Product struct {
 		Attributes struct {
-			Group     string `json:"group"`
-			Operation string `json:"operation"`
+			Group              string `json:"group"`
+			Operation          string `json:"operation"`
+			GroupDescription   string `json:"groupDescription"`
+			RequestDescription string `json:"requestDescription"`
+			InstanceType       string `json:"instanceType"`
+			InstanceTypeFamily string `json:"instanceTypeFamily"`
 		} `json:"attributes"`
 	} `json:"product"`
 	Terms struct {
@@ -47,7 +51,7 @@ func GetCost(pds []PriceDimensions, v float64) float64 {
 	return total
 }
 
-func GetPriceMap(pricingOutput *pricing.GetProductsOutput) (map[string][]PriceDimensions, error) {
+func GetPriceMap(pricingOutput *pricing.GetProductsOutput, field string) (map[string][]PriceDimensions, error) {
 	priceMap := make(map[string][]PriceDimensions)
 
 	if pricingOutput != nil && len(pricingOutput.PriceList) > 0 {
@@ -58,15 +62,30 @@ func GetPriceMap(pricingOutput *pricing.GetProductsOutput) (map[string][]PriceDi
 				return nil, fmt.Errorf("failed to unmarshal JSON: %w", err)
 			}
 
-			group := price.Product.Attributes.Group
+			var key string
+			switch field {
+			case "group":
+				key = price.Product.Attributes.Group
+			case "operation":
+				key = price.Product.Attributes.Operation
+			case "groupDescription":
+				key = price.Product.Attributes.GroupDescription
+			case "requestDescription":
+				key = price.Product.Attributes.RequestDescription
+			case "instanceType":
+				key = price.Product.Attributes.InstanceType
+			case "instanceTypeFamily":
+				key = price.Product.Attributes.InstanceTypeFamily
+			}
+
 			unitPrices := []PriceDimensions{}
 			for _, pd := range price.Terms.OnDemand {
 				for _, p := range pd.PriceDimensions {
 					unitPrices = append(unitPrices, p)
 				}
 			}
 
-			priceMap[group] = unitPrices
+			priceMap[key] = unitPrices
 		}
 	}
 

diff --git a/providers/aws/utils/utils_test.go b/providers/aws/utils/utils_test.go
@@ -49,6 +49,7 @@ func TestGetCost(t *testing.T) {
 func TestGetPriceMap(t *testing.T) {
 	testCases := []struct {
 		inputPriceList       []string
+		field                string
 		expectedNumProducts  int
 		expectedNumPriceDims map[string]int
 	}{
@@ -77,6 +78,7 @@ func TestGetPriceMap(t *testing.T) {
 						}
 					}
 				}`},
+			field:                "group",
 			expectedNumProducts:  1,
 			expectedNumPriceDims: map[string]int{"TestGroup": 1},
 		},
@@ -153,16 +155,46 @@ func TestGetPriceMap(t *testing.T) {
 					}
 				}`,
 			},
+			field:                "group",
 			expectedNumProducts:  2,
 			expectedNumPriceDims: map[string]int{"TestGroup1": 2, "TestGroup2": 3},
 		},
+		// Minimal valid JSON input with a single product, one price dimension & "instanceType" attribute
+		{
+			inputPriceList: []string{`
+				{
+					"product": {
+						"attributes": {
+							"instanceType": "TestInstanceType"
+						}
+					},
+					"terms": {
+						"OnDemand": {
+							"test_term": {
+								"priceDimensions": {
+									"test_price_dimension": {
+										"beginRange": "0",
+										"endRange": "Inf",
+										"pricePerUnit": {
+											"USD": "0.1"
+										}
+									}
+								}
+							}
+						}
+					}
+				}`},
+			field:                "instanceType",
+			expectedNumProducts:  1,
+			expectedNumPriceDims: map[string]int{"TestInstanceType": 1},
+		},
 	}
 	for i, testCase := range testCases {
 		t.Run(fmt.Sprintf("Test case %d", i+1), func(t *testing.T) {
 			output := pricing.GetProductsOutput{
 				PriceList: testCase.inputPriceList,
 			}
-			priceMap, err := GetPriceMap(&output)
+			priceMap, err := GetPriceMap(&output, "group")
 			if err != nil {
 				t.Errorf("Expected no error, but got: %v", err)
 			}
@@ -186,15 +218,15 @@ func TestGetPriceMap_InvalidJSON(t *testing.T) {
 	output := pricing.GetProductsOutput{
 		PriceList: []string{invalidJSON},
 	}
-	_, err := GetPriceMap(&output)
+	_, err := GetPriceMap(&output, "group")
 	if err == nil {
 		t.Error("Expected an error, but got nil")
 	}
 }
 
 func TestGetPriceMap_NoPricingOutput(t *testing.T) {
 	// PricingOutput is nil
-	priceMap, err := GetPriceMap(nil)
+	priceMap, err := GetPriceMap(nil, "group")
 	if err != nil {
 		t.Errorf("Expected no error, but got: %v", err)
 	}