github.com/brownsys/tracing-framework-go@v0.0.0-20161210174012-0542a62412fe/other/trace/baggage/baggage_test.go (about)

     1  package baggage
     2  
     3  import (
     4  	"encoding/binary"
     5  	"fmt"
     6  	"reflect"
     7  	"testing"
     8  
     9  	"github.com/golang/protobuf/proto"
    10  
    11  	lproto "github.com/brownsys/tracing-framework-go/trace/baggage/internal/proto"
    12  )
    13  
    14  func TestUnmarshalBasic(t *testing.T) {
    15  	testCaseUnmarshal(t, nil, ByteNamespaces{})
    16  	testCaseUnmarshal(t, []testNamespace{{}}, ByteNamespaces{"": ByteBaggage{}})
    17  	testCaseUnmarshal(t, []testNamespace{{key: []byte("foo"),
    18  		vals: []testBag{{key: []byte("bar"), vals: [][]byte{[]byte("baz")}}}}},
    19  		ByteNamespaces{"foo": ByteBaggage{"bar": [][]byte{[]byte("baz")}}})
    20  }
    21  
    22  type unmarshalableUint32 uint32
    23  
    24  func (u *unmarshalableUint32) UnmarshalBaggage(buf []byte) error {
    25  	if len(buf) != 4 {
    26  		return fmt.Errorf("invalid buffer length: %v", len(buf))
    27  	}
    28  	*u = unmarshalableUint32(binary.BigEndian.Uint32(buf))
    29  	return nil
    30  }
    31  
    32  func TestUnmarshalInterface(t *testing.T) {
    33  	type tt1 map[string][]unmarshalableUint32
    34  	type t1 map[string]tt1
    35  	testCaseUnmarshal(t, nil, t1{})
    36  	testCaseUnmarshal(t, []testNamespace{{}}, t1{"": tt1{}})
    37  	testCaseUnmarshal(t, []testNamespace{{key: []byte("foo"),
    38  		vals: []testBag{{key: []byte("bar"), vals: [][]byte{[]byte{1, 2, 3, 4}}}}}},
    39  		t1{"foo": tt1{"bar": []unmarshalableUint32{0x1020304}}})
    40  
    41  	type tt2 map[unmarshalableUint32][][]byte
    42  	type t2 map[string]tt2
    43  	testCaseUnmarshal(t, nil, t2{})
    44  	testCaseUnmarshal(t, []testNamespace{{}}, t2{"": tt2{}})
    45  	testCaseUnmarshal(t, []testNamespace{{key: []byte("foo"),
    46  		vals: []testBag{{key: []byte{0, 1, 2, 3}, vals: [][]byte{[]byte("bar")}},
    47  			{key: []byte{0, 1, 2, 3}, vals: [][]byte{[]byte("baz")}}}}},
    48  		t2{"foo": tt2{0x10203: [][]byte{[]byte("baz")}}})
    49  
    50  	type tt3 map[string][][]byte
    51  	type t3 map[unmarshalableUint32]tt3
    52  	testCaseUnmarshal(t, nil, t3{})
    53  	testCaseUnmarshal(t, []testNamespace{{key: []byte{0, 1, 2, 3}, vals: []testBag{}},
    54  		{key: []byte{0, 1, 2, 3}, vals: []testBag{{key: []byte("foo"), vals: [][]byte{[]byte("bar")}}}}},
    55  		t3{0x10203: tt3{"foo": [][]byte{[]byte("bar")}}})
    56  }
    57  
    58  // expect must be a valid type (map[T]map[U][]V)
    59  func testCaseUnmarshal(t *testing.T, ns []testNamespace, expect interface{}) {
    60  	msg := testNamespaceToProto(ns)
    61  	buf, err := proto.Marshal(&msg)
    62  	if err != nil {
    63  		t.Fatalf("unexpected error: %v", err)
    64  	}
    65  
    66  	typ := reflect.TypeOf(expect)
    67  	got := reflect.New(typ)
    68  	got.Elem().Set(reflect.MakeMap(typ))
    69  	err = Unmarshal(buf, got.Interface())
    70  	if err != nil {
    71  		t.Errorf("unexpected error: %v", err)
    72  		return
    73  	}
    74  
    75  	if !reflect.DeepEqual(got.Elem().Interface(), expect) {
    76  		t.Errorf("unexpected result: got %#v; want %#v", got.Elem().Interface(), expect)
    77  	}
    78  }
    79  
    80  type testNamespace struct {
    81  	key  []byte
    82  	vals []testBag
    83  }
    84  
    85  type testBag struct {
    86  	key  []byte
    87  	vals [][]byte
    88  }
    89  
    90  func testNamespaceToProto(t []testNamespace) lproto.BaggageMessage {
    91  	var msg lproto.BaggageMessage
    92  	for _, ns := range t {
    93  		var pns lproto.BaggageMessage_NamespaceData
    94  		// always allocate even if ns.key == nil
    95  		pns.Key = append([]byte{}, ns.key...)
    96  		for _, bag := range ns.vals {
    97  			var pbag lproto.BaggageMessage_BagData
    98  			// always allocate even if bag.key == nil
    99  			pbag.Key = append([]byte{}, bag.key...)
   100  			pbag.Value = bag.vals
   101  			pns.Bag = append(pns.Bag, &pbag)
   102  		}
   103  		msg.Namespace = append(msg.Namespace, &pns)
   104  	}
   105  	return msg
   106  }