passing context and refactored to table driven utests

Signed-off-by: pmahindrakar-oss <[email protected]>
flyteorg · Sep 18, 2024 · c140b8e · c140b8e
1 parent fceba3c
commit c140b8e
Show file tree

Hide file tree

Showing 3 changed files with 113 additions and 74 deletions.
diff --git a/flytepropeller/pkg/controller/config/config.go b/flytepropeller/pkg/controller/config/config.go
@@ -191,24 +191,24 @@ type LiteralOffloadingConfig struct {
 }
 
 // IsSupportedSDKVersion returns true if the provided SDK and version are supported by the literal offloading config.
-func (l LiteralOffloadingConfig) IsSupportedSDKVersion(sdk string, versionString string) bool {
+func (l LiteralOffloadingConfig) IsSupportedSDKVersion(ctx context.Context, sdk string, versionString string) bool {
 	regexMatches := sanitizeProtoRegex.FindStringSubmatch(versionString)
 	if len(regexMatches) > 1 {
-		logger.Infof(context.TODO(), "original: %s, semVer: %s", versionString, regexMatches[1])
+		logger.Infof(ctx, "original: %s, semVer: %s", versionString, regexMatches[1])
 	} else {
-		logger.Warnf(context.TODO(), "no match found for: %s", versionString)
+		logger.Warnf(ctx, "no match found for: %s", versionString)
 		return false
 	}
 	version, err := semver.NewVersion(regexMatches[1])
 	if err != nil {
-		logger.Warnf(context.TODO(), "Failed to parse version %s", versionString)
+		logger.Warnf(ctx, "Failed to parse version %s", versionString)
 		return false
 	}
 	if leastSupportedVersion, ok := l.SupportedSDKVersions[sdk]; ok {
 		c, err := semver.NewConstraint(fmt.Sprintf(">= %s", leastSupportedVersion))
 		if err != nil {
 			// This should never happen
-			logger.Warnf(context.TODO(), "Failed to parse version constraint %s", leastSupportedVersion)
+			logger.Warnf(ctx, "Failed to parse version constraint %s", leastSupportedVersion)
 			return false
 		}
 		return c.Check(version)

diff --git a/flytepropeller/pkg/controller/config/config_test.go b/flytepropeller/pkg/controller/config/config_test.go
@@ -1,87 +1,126 @@
 package config
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestIsSupportedSDKVersion(t *testing.T) {
-	t.Run("supported version", func(t *testing.T) {
-		config := LiteralOffloadingConfig{
-			SupportedSDKVersions: map[string]string{
-				"flytekit": "0.16.0",
+	ctx := context.Background()
+	tests := []struct {
+		name           string
+		config         LiteralOffloadingConfig
+		sdk            string
+		version        string
+		expectedResult bool
+	}{
+		{
+			name: "supported version",
+			config: LiteralOffloadingConfig{
+				SupportedSDKVersions: map[string]string{
+					"flytekit": "0.16.0",
+				},
 			},
-		}
-		assert.True(t, config.IsSupportedSDKVersion("flytekit", "0.16.0"))
-	})
-
-	t.Run("unsupported version", func(t *testing.T) {
-		config := LiteralOffloadingConfig{
-			SupportedSDKVersions: map[string]string{
-				"flytekit": "0.16.0",
+			sdk:            "flytekit",
+			version:        "0.16.0",
+			expectedResult: true,
+		},
+		{
+			name: "unsupported version",
+			config: LiteralOffloadingConfig{
+				SupportedSDKVersions: map[string]string{
+					"flytekit": "0.16.0",
+				},
 			},
-		}
-		assert.False(t, config.IsSupportedSDKVersion("flytekit", "0.15.0"))
-	})
-
-	t.Run("unsupported SDK", func(t *testing.T) {
-		config := LiteralOffloadingConfig{
-			SupportedSDKVersions: map[string]string{
-				"flytekit": "0.16.0",
+			sdk:            "flytekit",
+			version:        "0.15.0",
+			expectedResult: false,
+		},
+		{
+			name: "unsupported SDK",
+			config: LiteralOffloadingConfig{
+				SupportedSDKVersions: map[string]string{
+					"flytekit": "0.16.0",
+				},
 			},
-		}
-		assert.False(t, config.IsSupportedSDKVersion("unknown", "0.16.0"))
-	})
-
-	t.Run("invalid version", func(t *testing.T) {
-		config := LiteralOffloadingConfig{
-			SupportedSDKVersions: map[string]string{
-				"flytekit": "0.16.0",
+			sdk:            "unknown",
+			version:        "0.16.0",
+			expectedResult: false,
+		},
+		{
+			name: "invalid version",
+			config: LiteralOffloadingConfig{
+				SupportedSDKVersions: map[string]string{
+					"flytekit": "0.16.0",
+				},
 			},
-		}
-		assert.False(t, config.IsSupportedSDKVersion("flytekit", "invalid"))
-	})
-
-	t.Run("invalid constraint", func(t *testing.T) {
-		config := LiteralOffloadingConfig{
-			SupportedSDKVersions: map[string]string{
-				"flytekit": "invalid",
+			sdk:            "flytekit",
+			version:        "invalid",
+			expectedResult: false,
+		},
+		{
+			name: "invalid constraint",
+			config: LiteralOffloadingConfig{
+				SupportedSDKVersions: map[string]string{
+					"flytekit": "invalid",
+				},
 			},
-		}
-		assert.False(t, config.IsSupportedSDKVersion("flytekit", "0.16.0"))
-	})
-
-	t.Run("supported dev version", func(t *testing.T) {
-		config := LiteralOffloadingConfig{
-			SupportedSDKVersions: map[string]string{
-				"flytekit": "1.13.4",
+			sdk:            "flytekit",
+			version:        "0.16.0",
+			expectedResult: false,
+		},
+		{
+			name: "supported dev version",
+			config: LiteralOffloadingConfig{
+				SupportedSDKVersions: map[string]string{
+					"flytekit": "1.13.4",
+				},
 			},
-		}
-		assert.True(t, config.IsSupportedSDKVersion("flytekit", "1.13.4.dev12+g990b450ea.d20240917"))
-	})
-	t.Run("supported beta version", func(t *testing.T) {
-		config := LiteralOffloadingConfig{
-			SupportedSDKVersions: map[string]string{
-				"flytekit": "1.13.4",
+			sdk:            "flytekit",
+			version:        "1.13.4.dev12+g990b450ea.d20240917",
+			expectedResult: true,
+		},
+		{
+			name: "supported beta version",
+			config: LiteralOffloadingConfig{
+				SupportedSDKVersions: map[string]string{
+					"flytekit": "1.13.4",
+				},
 			},
-		}
-		assert.True(t, config.IsSupportedSDKVersion("flytekit", "v1.13.6b0"))
-	})
-	t.Run("unsupported dev version", func(t *testing.T) {
-		config := LiteralOffloadingConfig{
-			SupportedSDKVersions: map[string]string{
-				"flytekit": "1.13.4",
+			sdk:            "flytekit",
+			version:        "v1.13.6b0",
+			expectedResult: true,
+		},
+		{
+			name: "unsupported dev version",
+			config: LiteralOffloadingConfig{
+				SupportedSDKVersions: map[string]string{
+					"flytekit": "1.13.4",
+				},
 			},
-		}
-		assert.False(t, config.IsSupportedSDKVersion("flytekit", "1.13.3.dev12+g990b450ea.d20240917"))
-	})
-	t.Run("unsupported beta version", func(t *testing.T) {
-		config := LiteralOffloadingConfig{
-			SupportedSDKVersions: map[string]string{
-				"flytekit": "1.13.4",
+			sdk:            "flytekit",
+			version:        "1.13.3.dev12+g990b450ea.d20240917",
+			expectedResult: false,
+		},
+		{
+			name: "unsupported beta version",
+			config: LiteralOffloadingConfig{
+				SupportedSDKVersions: map[string]string{
+					"flytekit": "1.13.4",
+				},
 			},
-		}
-		assert.False(t, config.IsSupportedSDKVersion("flytekit", "v1.13.3b0"))
-	})
+			sdk:            "flytekit",
+			version:        "v1.13.3b0",
+			expectedResult: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := tt.config.IsSupportedSDKVersion(ctx, tt.sdk, tt.version)
+			assert.Equal(t, tt.expectedResult, result)
+		})
+	}
 }
diff --git a/flytepropeller/pkg/controller/nodes/common/utils.go b/flytepropeller/pkg/controller/nodes/common/utils.go
@@ -147,7 +147,7 @@ func CheckOffloadingCompat(ctx context.Context, nCtx interfaces.NodeExecutionCon
 		return &phaseInfo
 	}
 	runtimeData := taskNode.CoreTask().GetMetadata().GetRuntime()
-	if !literalOffloadingConfig.IsSupportedSDKVersion(runtimeData.GetType().String(), runtimeData.GetVersion()) {
+	if !literalOffloadingConfig.IsSupportedSDKVersion(ctx, runtimeData.GetType().String(), runtimeData.GetVersion()) {
 		if !literalOffloadingConfig.Enabled {
 			errMsg := fmt.Sprintf("task [%s] is trying to consume offloaded literals but feature is not enabled", taskID)
 			logger.Errorf(ctx, errMsg)