lib/go/thrift/header_context_test.go

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied. See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package thrift

import (
	"context"
	"reflect"
	"testing"
)

func TestSetGetUnsetHeader(t *testing.T) {
	const (
		key   = "foo"
		value = "bar"
	)
	ctx := context.Background()

	ctx = SetHeader(ctx, key, value)

	checkGet := func(t *testing.T, ctx context.Context) {
		t.Helper()
		got, ok := GetHeader(ctx, key)
		if !ok {
			t.Fatalf("Cannot get header %q back after setting it.", key)
		}
		if got != value {
			t.Fatalf("Header value expected %q, got %q instead", value, got)
		}
	}

	checkGet(t, ctx)

	t.Run(
		"NoConflicts",
		func(t *testing.T) {
			type otherType string
			const otherValue = "bar2"

			ctx = context.WithValue(ctx, otherType(key), otherValue)
			checkGet(t, ctx)
		},
	)

	t.Run(
		"GetHeaderOnNonExistKey",
		func(t *testing.T) {
			const otherKey = "foo2"

			if _, ok := GetHeader(ctx, otherKey); ok {
				t.Errorf("GetHeader returned ok on non-existing key %q", otherKey)
			}
		},
	)

	t.Run(
		"Unset",
		func(t *testing.T) {
			ctx := UnsetHeader(ctx, key)

			if _, ok := GetHeader(ctx, key); ok {
				t.Errorf("GetHeader returned ok on unset key %q", key)
			}
		},
	)
}

func TestReadKeyList(t *testing.T) {
	headers := THeaderMap{
		"key1": "value1",
		"key2": "value2",
	}
	ctx := context.Background()

	ctx = AddReadTHeaderToContext(ctx, headers)

	got := make(THeaderMap)
	keys := GetReadHeaderList(ctx)
	t.Logf("keys: %+v", keys)
	for _, key := range keys {
		value, ok := GetHeader(ctx, key)
		if ok {
			got[key] = value
		} else {
			t.Errorf("Cannot get key %q from context", key)
		}
	}

	if !reflect.DeepEqual(headers, got) {
		t.Errorf("Expected header map %+v, got %+v", headers, got)
	}

	writtenKeys := GetWriteHeaderList(ctx)
	if len(writtenKeys) > 0 {
		t.Errorf(
			"Expected empty GetWriteHeaderList() result, got %+v",
			writtenKeys,
		)
	}
}

func TestWriteKeyList(t *testing.T) {
	keys := []string{
		"key1",
		"key2",
	}
	ctx := context.Background()

	ctx = SetWriteHeaderList(ctx, keys)
	got := GetWriteHeaderList(ctx)

	if !reflect.DeepEqual(keys, got) {
		t.Errorf("Expected header keys %+v, got %+v", keys, got)
	}

	readKeys := GetReadHeaderList(ctx)
	if len(readKeys) > 0 {
		t.Errorf(
			"Expected empty GetReadHeaderList() result, got %+v",
			readKeys,
		)
	}
}