Merge pull request #368 from Guojunzhou-git/main

param for NewMsgGrpc and NewSagaGrpc to setup BranchHeaders

client/dtmgrpc/msg.go

@@ -24,8 +24,14 @@ type MsgGrpc struct {
 }
 
 // NewMsgGrpc create new msg
-func NewMsgGrpc(server string, gid string) *MsgGrpc {
-	return &MsgGrpc{Msg: *dtmcli.NewMsg(server, gid)}
+func NewMsgGrpc(server string, gid string, opts ...TransBaseOption) *MsgGrpc {
+	mg := &MsgGrpc{Msg: *dtmcli.NewMsg(server, gid)}
+
+	for _, opt := range opts {
+		opt(&mg.TransBase)
+	}
+
+	return mg
 }
 
 // Add add a new step

client/dtmgrpc/options.go

@@ -0,0 +1,15 @@
+package dtmgrpc
+
+import (
+	"github.com/dtm-labs/dtm/client/dtmcli/dtmimp"
+)
+
+// TransBaseOption setup func for TransBase
+type TransBaseOption func(tb *dtmimp.TransBase)
+
+// WithBranchHeaders setup TransBase.BranchHeaders
+func WithBranchHeaders(headers map[string]string) TransBaseOption {
+	return func(tb *dtmimp.TransBase) {
+		tb.BranchHeaders = headers
+	}
+}

client/dtmgrpc/options_test.go

@@ -0,0 +1,104 @@
+package dtmgrpc
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/dtm-labs/dtm/client/dtmcli"
+)
+
+// TestNewMsgGrpc ut for NewMsgGrpc
+func TestNewMsgGrpc(t *testing.T) {
+	var (
+		server            = "dmt_server_address"
+		gidNoOptions      = "msg_no_setup_options"
+		gidTraceIDXXX     = "msg_setup_options_trace_id_xxx"
+		msgWithTraceIDXXX = &MsgGrpc{Msg: *dtmcli.NewMsg(server, gidTraceIDXXX)}
+		traceIDHeaders    = map[string]string{
+			"x-trace-id": "xxx",
+		}
+	)
+	msgWithTraceIDXXX.BranchHeaders = traceIDHeaders
+	type args struct {
+		gid  string
+		opts []TransBaseOption
+	}
+	tests := []struct {
+		name string
+		args args
+		want *MsgGrpc
+	}{
+		{
+			name: "no setup options",
+			args: args{gid: gidNoOptions},
+			want: &MsgGrpc{Msg: *dtmcli.NewMsg(server, gidNoOptions)},
+		},
+		{
+			name: "msg with trace_id",
+			args: args{
+				gid: gidTraceIDXXX,
+				opts: []TransBaseOption{
+					WithBranchHeaders(traceIDHeaders),
+				},
+			},
+			want: msgWithTraceIDXXX,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := NewMsgGrpc(server, tt.args.gid, tt.args.opts...)
+			t.Logf("TestNewMsgGrpc %s got %+v\n", tt.name, got)
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("NewMsgGrpc() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+// TestNewSagaGrpc ut for NewSagaGrpc
+func TestNewSagaGrpc(t *testing.T) {
+	var (
+		server             = "dmt_server_address"
+		gidNoOptions       = "msg_no_setup_options"
+		gidTraceIDXXX      = "msg_setup_options_trace_id_xxx"
+		sagaWithTraceIDXXX = &SagaGrpc{Saga: *dtmcli.NewSaga(server, gidTraceIDXXX)}
+		traceIDHeaders     = map[string]string{
+			"x-trace-id": "xxx",
+		}
+	)
+	sagaWithTraceIDXXX.BranchHeaders = traceIDHeaders
+	type args struct {
+		gid  string
+		opts []TransBaseOption
+	}
+	tests := []struct {
+		name string
+		args args
+		want *SagaGrpc
+	}{
+		{
+			name: "no setup options",
+			args: args{gid: gidNoOptions},
+			want: &SagaGrpc{Saga: *dtmcli.NewSaga(server, gidNoOptions)},
+		},
+		{
+			name: "msg with trace_id",
+			args: args{
+				gid: gidTraceIDXXX,
+				opts: []TransBaseOption{
+					WithBranchHeaders(traceIDHeaders),
+				},
+			},
+			want: sagaWithTraceIDXXX,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := NewSagaGrpc(server, tt.args.gid, tt.args.opts...)
+			t.Logf("TestNewSagaGrpc %s got %+v\n", tt.name, got)
+			if !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("NewSagaGrpc() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

client/dtmgrpc/saga.go

@@ -18,8 +18,14 @@ type SagaGrpc struct {
 }
 
 // NewSagaGrpc create a saga
-func NewSagaGrpc(server string, gid string) *SagaGrpc {
-	return &SagaGrpc{Saga: *dtmcli.NewSaga(server, gid)}
+func NewSagaGrpc(server string, gid string, opts ...TransBaseOption) *SagaGrpc {
+	sg := &SagaGrpc{Saga: *dtmcli.NewSaga(server, gid)}
+
+	for _, opt := range opts {
+		opt(&sg.TransBase)
+	}
+
+	return sg
 }
 
 // Add add a saga step