Selaa lähdekoodia

Add configuration toggle to disable AWS spot data feed (#3657)

Kush Agarwal 1 kuukausi sitten
vanhempi
sitoutus
302e48ec8e
3 muutettua tiedostoa jossa 175 lisäystä ja 2 poistoa
  1. 22 2
      pkg/cloud/aws/provider.go
  2. 152 0
      pkg/cloud/aws/provider_test.go
  3. 1 0
      pkg/cloud/models/models.go

+ 22 - 2
pkg/cloud/aws/provider.go

@@ -855,8 +855,28 @@ func (aws *AWS) getRegionPricing(nodeList []*clustercache.Node) (*http.Response,
 
 // SpotRefreshEnabled determines whether the required configs to run the spot feed query have been set up
 func (aws *AWS) SpotRefreshEnabled() bool {
-	// Need a valid value for at least one of these fields to consider spot pricing as enabled
-	return len(aws.SpotDataBucket) != 0 || len(aws.SpotDataRegion) != 0 || len(aws.ProjectID) != 0
+	// Guard against nil receiver
+	if aws == nil {
+		return false
+	}
+
+	// Fallback if config is not initialized
+	if aws.Config == nil {
+		return len(aws.SpotDataBucket) != 0 ||
+			len(aws.SpotDataRegion) != 0 ||
+			len(aws.ProjectID) != 0
+	}
+
+	// Check if spot data feed is explicitly disabled via config
+	c, err := aws.Config.GetCustomPricingData()
+	if err == nil && c.SpotDataFeedEnabled == "false" {
+		return false
+	}
+
+	// Default behavior
+	return len(aws.SpotDataBucket) != 0 ||
+		len(aws.SpotDataRegion) != 0 ||
+		len(aws.ProjectID) != 0
 }
 
 // DownloadPricingData fetches data from the AWS Pricing API

+ 152 - 0
pkg/cloud/aws/provider_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/opencost/opencost/core/pkg/clustercache"
 	"github.com/opencost/opencost/pkg/cloud/models"
+	"github.com/opencost/opencost/pkg/config"
 	v1 "k8s.io/api/core/v1"
 )
 
@@ -841,3 +842,154 @@ func TestAWS_getFargatePod(t *testing.T) {
 		})
 	}
 }
+
+// fakeProviderConfig implements models.ProviderConfig for testing
+type fakeProviderConfig struct {
+	customPricing *models.CustomPricing
+}
+
+func (f *fakeProviderConfig) GetCustomPricingData() (*models.CustomPricing, error) {
+	if f.customPricing != nil {
+		return f.customPricing, nil
+	}
+	return &models.CustomPricing{}, nil
+}
+
+func (f *fakeProviderConfig) Update(func(*models.CustomPricing) error) (*models.CustomPricing, error) {
+	return f.customPricing, nil
+}
+
+func (f *fakeProviderConfig) UpdateFromMap(map[string]string) (*models.CustomPricing, error) {
+	return f.customPricing, nil
+}
+
+func (f *fakeProviderConfig) ConfigFileManager() *config.ConfigFileManager {
+	return nil
+}
+
+func TestAWS_SpotRefreshEnabled(t *testing.T) {
+	tests := []struct {
+		name                string
+		spotDataBucket      string
+		spotDataRegion      string
+		projectID           string
+		spotDataFeedEnabled string
+		want                bool
+	}{
+		{
+			name:                "disabled via config - with bucket",
+			spotDataBucket:      "my-bucket",
+			spotDataRegion:      "us-east-1",
+			projectID:           "123456789",
+			spotDataFeedEnabled: "false",
+			want:                false,
+		},
+		{
+			name:                "disabled via config - with projectID only",
+			projectID:           "123456789",
+			spotDataFeedEnabled: "false",
+			want:                false,
+		},
+		{
+			name:                "enabled by default - with bucket",
+			spotDataBucket:      "my-bucket",
+			spotDataRegion:      "us-east-1",
+			projectID:           "123456789",
+			spotDataFeedEnabled: "",
+			want:                true,
+		},
+		{
+			name:                "enabled explicitly - with bucket",
+			spotDataBucket:      "my-bucket",
+			spotDataRegion:      "us-east-1",
+			projectID:           "123456789",
+			spotDataFeedEnabled: "true",
+			want:                true,
+		},
+		{
+			name:                "no spot config - disabled",
+			spotDataBucket:      "",
+			spotDataRegion:      "",
+			projectID:           "",
+			spotDataFeedEnabled: "",
+			want:                false,
+		},
+		{
+			name:                "no spot config - but explicitly enabled",
+			spotDataBucket:      "",
+			spotDataRegion:      "",
+			projectID:           "",
+			spotDataFeedEnabled: "true",
+			want:                false,
+		},
+		{
+			name:                "only projectID set - enabled by default",
+			projectID:           "123456789",
+			spotDataFeedEnabled: "",
+			want:                true,
+		},
+		{
+			name:                "only bucket set - enabled by default",
+			spotDataBucket:      "my-bucket",
+			spotDataFeedEnabled: "",
+			want:                true,
+		},
+		{
+			name:                "only region set - enabled by default",
+			spotDataRegion:      "us-east-1",
+			spotDataFeedEnabled: "",
+			want:                true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			aws := &AWS{
+				SpotDataBucket: tt.spotDataBucket,
+				SpotDataRegion: tt.spotDataRegion,
+				ProjectID:      tt.projectID,
+				Config: &fakeProviderConfig{
+					customPricing: &models.CustomPricing{
+						SpotDataFeedEnabled: tt.spotDataFeedEnabled,
+					},
+				},
+			}
+
+			got := aws.SpotRefreshEnabled()
+			if got != tt.want {
+				t.Errorf("AWS.SpotRefreshEnabled() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+
+	// Test nil Config scenario to ensure no panic
+	t.Run("nil config - falls back to field check", func(t *testing.T) {
+		aws := &AWS{
+			SpotDataBucket: "my-bucket",
+			SpotDataRegion: "us-east-1",
+			ProjectID:      "123456789",
+			Config:         nil, // nil Config should not cause panic
+		}
+
+		got := aws.SpotRefreshEnabled()
+		want := true // Should fall back to field-based check
+		if got != want {
+			t.Errorf("AWS.SpotRefreshEnabled() with nil Config = %v, want %v", got, want)
+		}
+	})
+
+	t.Run("nil config - no spot fields", func(t *testing.T) {
+		aws := &AWS{
+			SpotDataBucket: "",
+			SpotDataRegion: "",
+			ProjectID:      "",
+			Config:         nil, // nil Config should not cause panic
+		}
+
+		got := aws.SpotRefreshEnabled()
+		want := false // No fields set, should return false
+		if got != want {
+			t.Errorf("AWS.SpotRefreshEnabled() with nil Config and no fields = %v, want %v", got, want)
+		}
+	})
+}

+ 1 - 0
pkg/cloud/models/models.go

@@ -154,6 +154,7 @@ type CustomPricing struct {
 	AwsSpotDataRegion            string `json:"awsSpotDataRegion,omitempty"`
 	AwsSpotDataBucket            string `json:"awsSpotDataBucket,omitempty"`
 	AwsSpotDataPrefix            string `json:"awsSpotDataPrefix,omitempty"`
+	SpotDataFeedEnabled          string `json:"spotDataFeedEnabled,omitempty"`
 	ProjectID                    string `json:"projectID,omitempty"`
 	AthenaProjectID              string `json:"athenaProjectID,omitempty"`
 	AthenaBucketName             string `json:"athenaBucketName"`