Merge pull request #215 from projectdiscovery/feat-ctx-or-default

Adding value or default context helper
projectdiscovery · Aug 9, 2023 · cc55449 · cc55449
2 parents 043433a + 6a90183
commit cc55449
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/context/context.go b/context/context.go
@@ -7,8 +7,12 @@ import (
 
 var ErrIncorrectNumberOfItems = errors.New("number of items is not even")
 
+var DefaultContext = context.TODO()
+
+type ContextArg string
+
 // WithValues combines multiple key-value into an existing context
-func WithValues(ctx context.Context, keyValue ...string) (context.Context, error) {
+func WithValues(ctx context.Context, keyValue ...ContextArg) (context.Context, error) {
 	if len(keyValue)%2 != 0 {
 		return ctx, ErrIncorrectNumberOfItems
 	}
@@ -18,3 +22,12 @@ func WithValues(ctx context.Context, keyValue ...string) (context.Context, error
 	}
 	return ctx, nil
 }
+
+// ValueOrDefault returns default context if given is nil (using interface to avoid static check reporting)
+func ValueOrDefault(value interface{}) context.Context {
+	if ctx, ok := value.(context.Context); ok && ctx != nil {
+		return ctx
+	}
+
+	return DefaultContext
+}
diff --git a/context/context_test.go b/context/context_test.go
@@ -10,29 +10,29 @@ import (
 func TestWithValues(t *testing.T) {
 	type testCase struct {
 		name          string
-		keyValue      []string
+		keyValue      []ContextArg
 		expectedError error
-		expectedValue map[string]string
+		expectedValue map[ContextArg]ContextArg
 	}
 
 	var testCases = []testCase{
 		{
 			name:          "even number of key-value pairs",
-			keyValue:      []string{"key1", "value1", "key2", "value2"},
+			keyValue:      []ContextArg{"key1", "value1", "key2", "value2"},
 			expectedError: nil,
-			expectedValue: map[string]string{"key1": "value1", "key2": "value2"},
+			expectedValue: map[ContextArg]ContextArg{"key1": "value1", "key2": "value2"},
 		},
 		{
 			name:          "odd number of key-value pairs",
-			keyValue:      []string{"key1", "value1", "key2"},
+			keyValue:      []ContextArg{"key1", "value1", "key2"},
 			expectedError: ErrIncorrectNumberOfItems,
-			expectedValue: map[string]string{},
+			expectedValue: map[ContextArg]ContextArg{},
 		},
 		{
 			name:          "overwriting values",
-			keyValue:      []string{"key1", "value1", "key1", "newValue"},
+			keyValue:      []ContextArg{"key1", "value1", "key1", "newValue"},
 			expectedError: nil,
-			expectedValue: map[string]string{"key1": "newValue"},
+			expectedValue: map[ContextArg]ContextArg{"key1": "newValue"},
 		},
 	}
 	ctx := context.Background()