| @@ -28,3 +28,20 @@ func FromContext(ctx context.Context) (Metadata, bool) { | |||||||
| func NewContext(ctx context.Context, md Metadata) context.Context { | func NewContext(ctx context.Context, md Metadata) context.Context { | ||||||
| 	return context.WithValue(ctx, metaKey{}, md) | 	return context.WithValue(ctx, metaKey{}, md) | ||||||
| } | } | ||||||
|  |  | ||||||
|  | func AppendContext(ctx context.Context, patchMd Metadata, overwrite bool) context.Context { | ||||||
|  | 	md, _ := ctx.Value(metaKey{}).(Metadata) | ||||||
|  | 	cmd := make(Metadata) | ||||||
|  | 	for k, v := range md { | ||||||
|  | 		cmd[k] = v | ||||||
|  | 	} | ||||||
|  | 	for k, v := range patchMd { | ||||||
|  | 		if _, ok := cmd[k]; ok && !overwrite { | ||||||
|  | 			// skip | ||||||
|  | 		} else { | ||||||
|  | 			cmd[k] = v | ||||||
|  | 		} | ||||||
|  | 	} | ||||||
|  | 	return context.WithValue(ctx, metaKey{}, cmd) | ||||||
|  |  | ||||||
|  | } | ||||||
|   | |||||||
| @@ -2,6 +2,7 @@ package metadata | |||||||
|  |  | ||||||
| import ( | import ( | ||||||
| 	"context" | 	"context" | ||||||
|  | 	"reflect" | ||||||
| 	"testing" | 	"testing" | ||||||
| ) | ) | ||||||
|  |  | ||||||
| @@ -40,3 +41,42 @@ func TestMetadataContext(t *testing.T) { | |||||||
| 		t.Errorf("Expected metadata length 1 got %d", i) | 		t.Errorf("Expected metadata length 1 got %d", i) | ||||||
| 	} | 	} | ||||||
| } | } | ||||||
|  |  | ||||||
|  | func TestAppendContext(t *testing.T) { | ||||||
|  | 	type args struct { | ||||||
|  | 		existing  Metadata | ||||||
|  | 		append    Metadata | ||||||
|  | 		overwrite bool | ||||||
|  | 	} | ||||||
|  | 	tests := []struct { | ||||||
|  | 		name string | ||||||
|  | 		args args | ||||||
|  | 		want Metadata | ||||||
|  | 	}{ | ||||||
|  | 		{ | ||||||
|  | 			name: "matching key, overwrite false", | ||||||
|  | 			args: args{ | ||||||
|  | 				existing:  Metadata{"foo": "bar", "sumo": "demo"}, | ||||||
|  | 				append:    Metadata{"sumo": "demo2"}, | ||||||
|  | 				overwrite: false, | ||||||
|  | 			}, | ||||||
|  | 			want: Metadata{"foo": "bar", "sumo": "demo"}, | ||||||
|  | 		}, | ||||||
|  | 		{ | ||||||
|  | 			name: "matching key, overwrite true", | ||||||
|  | 			args: args{ | ||||||
|  | 				existing:  Metadata{"foo": "bar", "sumo": "demo"}, | ||||||
|  | 				append:    Metadata{"sumo": "demo2"}, | ||||||
|  | 				overwrite: true, | ||||||
|  | 			}, | ||||||
|  | 			want: Metadata{"foo": "bar", "sumo": "demo2"}, | ||||||
|  | 		}, | ||||||
|  | 	} | ||||||
|  | 	for _, tt := range tests { | ||||||
|  | 		t.Run(tt.name, func(t *testing.T) { | ||||||
|  | 			if got, _ := FromContext(AppendContext(NewContext(context.TODO(), tt.args.existing), tt.args.append, tt.args.overwrite)); !reflect.DeepEqual(got, tt.want) { | ||||||
|  | 				t.Errorf("AppendContext() = %v, want %v", got, tt.want) | ||||||
|  | 			} | ||||||
|  | 		}) | ||||||
|  | 	} | ||||||
|  | } | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user