diff --git a/client/dtmgrpc/msg.go b/client/dtmgrpc/msg.go index e08d2f3..e9b59da 100644 --- a/client/dtmgrpc/msg.go +++ b/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 diff --git a/client/dtmgrpc/options.go b/client/dtmgrpc/options.go new file mode 100644 index 0000000..1ea61c7 --- /dev/null +++ b/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 + } +} diff --git a/client/dtmgrpc/options_test.go b/client/dtmgrpc/options_test.go new file mode 100644 index 0000000..011c36b --- /dev/null +++ b/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) + } + }) + } +} diff --git a/client/dtmgrpc/saga.go b/client/dtmgrpc/saga.go index 9d801aa..cc13df8 100644 --- a/client/dtmgrpc/saga.go +++ b/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