diff --git a/backend/kitex_gen/coze/loop/observability/domain/trace/k-trace.go b/backend/kitex_gen/coze/loop/observability/domain/trace/k-trace.go index 4f0812316..9fcc649d2 100644 --- a/backend/kitex_gen/coze/loop/observability/domain/trace/k-trace.go +++ b/backend/kitex_gen/coze/loop/observability/domain/trace/k-trace.go @@ -203,8 +203,8 @@ func (p *TokenCost) FastRead(buf []byte) (int, error) { var l int var fieldTypeId thrift.TType var fieldId int16 - var issetInput bool = false - var issetOutput bool = false + var issetInputToken bool = false + var issetOutputToken bool = false for { fieldTypeId, fieldId, l, err = thrift.Binary.ReadFieldBegin(buf[offset:]) offset += l @@ -222,7 +222,7 @@ func (p *TokenCost) FastRead(buf []byte) (int, error) { if err != nil { goto ReadFieldError } - issetInput = true + issetInputToken = true } else { l, err = thrift.Binary.Skip(buf[offset:], fieldTypeId) offset += l @@ -237,7 +237,7 @@ func (p *TokenCost) FastRead(buf []byte) (int, error) { if err != nil { goto ReadFieldError } - issetOutput = true + issetOutputToken = true } else { l, err = thrift.Binary.Skip(buf[offset:], fieldTypeId) offset += l @@ -254,12 +254,12 @@ func (p *TokenCost) FastRead(buf []byte) (int, error) { } } - if !issetInput { + if !issetInputToken { fieldId = 1 goto RequiredFieldNotSetError } - if !issetOutput { + if !issetOutputToken { fieldId = 2 goto RequiredFieldNotSetError } @@ -284,7 +284,7 @@ func (p *TokenCost) FastReadField1(buf []byte) (int, error) { offset += l _field = v } - p.Input = _field + p.InputToken = _field return offset, nil } @@ -298,7 +298,7 @@ func (p *TokenCost) FastReadField2(buf []byte) (int, error) { offset += l _field = v } - p.Output = _field + p.OutputToken = _field return offset, nil } @@ -329,14 +329,14 @@ func (p *TokenCost) BLength() int { func (p *TokenCost) fastWriteField1(buf []byte, w thrift.NocopyWriter) int { offset := 0 offset += thrift.Binary.WriteFieldBegin(buf[offset:], thrift.I64, 1) - offset += thrift.Binary.WriteI64(buf[offset:], p.Input) + offset += thrift.Binary.WriteI64(buf[offset:], p.InputToken) return offset } func (p *TokenCost) fastWriteField2(buf []byte, w thrift.NocopyWriter) int { offset := 0 offset += thrift.Binary.WriteFieldBegin(buf[offset:], thrift.I64, 2) - offset += thrift.Binary.WriteI64(buf[offset:], p.Output) + offset += thrift.Binary.WriteI64(buf[offset:], p.OutputToken) return offset } @@ -360,9 +360,9 @@ func (p *TokenCost) DeepCopy(s interface{}) error { return fmt.Errorf("%T's type not matched %T", s, p) } - p.Input = src.Input + p.InputToken = src.InputToken - p.Output = src.Output + p.OutputToken = src.OutputToken return nil } diff --git a/backend/kitex_gen/coze/loop/observability/domain/trace/trace.go b/backend/kitex_gen/coze/loop/observability/domain/trace/trace.go index c9fc5c24a..89b63dc07 100644 --- a/backend/kitex_gen/coze/loop/observability/domain/trace/trace.go +++ b/backend/kitex_gen/coze/loop/observability/domain/trace/trace.go @@ -259,8 +259,8 @@ func (p *Trace) Field2DeepEqual(src *TokenCost) bool { } type TokenCost struct { - Input int64 `thrift:"input,1,required" frugal:"1,required,i64" json:"input" form:"input,required" query:"input,required"` - Output int64 `thrift:"output,2,required" frugal:"2,required,i64" json:"output" form:"output,required" query:"output,required"` + InputToken int64 `thrift:"input_token,1,required" frugal:"1,required,i64" json:"input_token" form:"input_token,required" query:"input_token,required"` + OutputToken int64 `thrift:"output_token,2,required" frugal:"2,required,i64" json:"output_token" form:"output_token,required" query:"output_token,required"` } func NewTokenCost() *TokenCost { @@ -270,36 +270,36 @@ func NewTokenCost() *TokenCost { func (p *TokenCost) InitDefault() { } -func (p *TokenCost) GetInput() (v int64) { +func (p *TokenCost) GetInputToken() (v int64) { if p != nil { - return p.Input + return p.InputToken } return } -func (p *TokenCost) GetOutput() (v int64) { +func (p *TokenCost) GetOutputToken() (v int64) { if p != nil { - return p.Output + return p.OutputToken } return } -func (p *TokenCost) SetInput(val int64) { - p.Input = val +func (p *TokenCost) SetInputToken(val int64) { + p.InputToken = val } -func (p *TokenCost) SetOutput(val int64) { - p.Output = val +func (p *TokenCost) SetOutputToken(val int64) { + p.OutputToken = val } var fieldIDToName_TokenCost = map[int16]string{ - 1: "input", - 2: "output", + 1: "input_token", + 2: "output_token", } func (p *TokenCost) Read(iprot thrift.TProtocol) (err error) { var fieldTypeId thrift.TType var fieldId int16 - var issetInput bool = false - var issetOutput bool = false + var issetInputToken bool = false + var issetOutputToken bool = false if _, err = iprot.ReadStructBegin(); err != nil { goto ReadStructBeginError @@ -320,7 +320,7 @@ func (p *TokenCost) Read(iprot thrift.TProtocol) (err error) { if err = p.ReadField1(iprot); err != nil { goto ReadFieldError } - issetInput = true + issetInputToken = true } else if err = iprot.Skip(fieldTypeId); err != nil { goto SkipFieldError } @@ -329,7 +329,7 @@ func (p *TokenCost) Read(iprot thrift.TProtocol) (err error) { if err = p.ReadField2(iprot); err != nil { goto ReadFieldError } - issetOutput = true + issetOutputToken = true } else if err = iprot.Skip(fieldTypeId); err != nil { goto SkipFieldError } @@ -346,12 +346,12 @@ func (p *TokenCost) Read(iprot thrift.TProtocol) (err error) { goto ReadStructEndError } - if !issetInput { + if !issetInputToken { fieldId = 1 goto RequiredFieldNotSetError } - if !issetOutput { + if !issetOutputToken { fieldId = 2 goto RequiredFieldNotSetError } @@ -381,7 +381,7 @@ func (p *TokenCost) ReadField1(iprot thrift.TProtocol) error { } else { _field = v } - p.Input = _field + p.InputToken = _field return nil } func (p *TokenCost) ReadField2(iprot thrift.TProtocol) error { @@ -392,7 +392,7 @@ func (p *TokenCost) ReadField2(iprot thrift.TProtocol) error { } else { _field = v } - p.Output = _field + p.OutputToken = _field return nil } @@ -429,10 +429,10 @@ WriteStructEndError: } func (p *TokenCost) writeField1(oprot thrift.TProtocol) (err error) { - if err = oprot.WriteFieldBegin("input", thrift.I64, 1); err != nil { + if err = oprot.WriteFieldBegin("input_token", thrift.I64, 1); err != nil { goto WriteFieldBeginError } - if err := oprot.WriteI64(p.Input); err != nil { + if err := oprot.WriteI64(p.InputToken); err != nil { return err } if err = oprot.WriteFieldEnd(); err != nil { @@ -445,10 +445,10 @@ WriteFieldEndError: return thrift.PrependError(fmt.Sprintf("%T write field 1 end error: ", p), err) } func (p *TokenCost) writeField2(oprot thrift.TProtocol) (err error) { - if err = oprot.WriteFieldBegin("output", thrift.I64, 2); err != nil { + if err = oprot.WriteFieldBegin("output_token", thrift.I64, 2); err != nil { goto WriteFieldBeginError } - if err := oprot.WriteI64(p.Output); err != nil { + if err := oprot.WriteI64(p.OutputToken); err != nil { return err } if err = oprot.WriteFieldEnd(); err != nil { @@ -475,10 +475,10 @@ func (p *TokenCost) DeepEqual(ano *TokenCost) bool { } else if p == nil || ano == nil { return false } - if !p.Field1DeepEqual(ano.Input) { + if !p.Field1DeepEqual(ano.InputToken) { return false } - if !p.Field2DeepEqual(ano.Output) { + if !p.Field2DeepEqual(ano.OutputToken) { return false } return true @@ -486,14 +486,14 @@ func (p *TokenCost) DeepEqual(ano *TokenCost) bool { func (p *TokenCost) Field1DeepEqual(src int64) bool { - if p.Input != src { + if p.InputToken != src { return false } return true } func (p *TokenCost) Field2DeepEqual(src int64) bool { - if p.Output != src { + if p.OutputToken != src { return false } return true diff --git a/backend/kitex_gen/coze/loop/prompt/openapi/coze.loop.prompt.openapi.go b/backend/kitex_gen/coze/loop/prompt/openapi/coze.loop.prompt.openapi.go index f0ff648fe..3abac5b11 100644 --- a/backend/kitex_gen/coze/loop/prompt/openapi/coze.loop.prompt.openapi.go +++ b/backend/kitex_gen/coze/loop/prompt/openapi/coze.loop.prompt.openapi.go @@ -1024,11 +1024,15 @@ func (p *PromptResultData) Field1DeepEqual(src []*PromptResult_) bool { } type ExecuteRequest struct { - WorkspaceID *int64 `thrift:"workspace_id,1,optional" frugal:"1,optional,i64" json:"workspace_id" form:"workspace_id" ` - PromptIdentifier *PromptQuery `thrift:"prompt_identifier,2,optional" frugal:"2,optional,PromptQuery" form:"prompt_identifier" json:"prompt_identifier,omitempty"` - VariableVals []*VariableVal `thrift:"variable_vals,10,optional" frugal:"10,optional,list" form:"variable_vals" json:"variable_vals,omitempty"` - Messages []*Message `thrift:"messages,11,optional" frugal:"11,optional,list" form:"messages" json:"messages,omitempty"` - Base *base.Base `thrift:"Base,255,optional" frugal:"255,optional,base.Base" form:"Base" json:"Base,omitempty" query:"Base"` + // 工作空间ID + WorkspaceID *int64 `thrift:"workspace_id,1,optional" frugal:"1,optional,i64" json:"workspace_id" form:"workspace_id" ` + // Prompt 标识 + PromptIdentifier *PromptQuery `thrift:"prompt_identifier,2,optional" frugal:"2,optional,PromptQuery" form:"prompt_identifier" json:"prompt_identifier,omitempty"` + // 变量值 + VariableVals []*VariableVal `thrift:"variable_vals,10,optional" frugal:"10,optional,list" form:"variable_vals" json:"variable_vals,omitempty"` + // 消息 + Messages []*Message `thrift:"messages,11,optional" frugal:"11,optional,list" form:"messages" json:"messages,omitempty"` + Base *base.Base `thrift:"Base,255,optional" frugal:"255,optional,base.Base" form:"Base" json:"Base,omitempty" query:"Base"` } func NewExecuteRequest() *ExecuteRequest { @@ -1935,9 +1939,12 @@ func (p *ExecuteResponse) Field255DeepEqual(src *base.BaseResp) bool { } type ExecuteData struct { - Message *Message `thrift:"message,1,optional" frugal:"1,optional,Message" form:"message" json:"message,omitempty" query:"message"` - FinishReason *string `thrift:"finish_reason,2,optional" frugal:"2,optional,string" form:"finish_reason" json:"finish_reason,omitempty" query:"finish_reason"` - Usage *TokenUsage `thrift:"usage,3,optional" frugal:"3,optional,TokenUsage" form:"usage" json:"usage,omitempty" query:"usage"` + // 消息 + Message *Message `thrift:"message,1,optional" frugal:"1,optional,Message" form:"message" json:"message,omitempty" query:"message"` + // 结束原因 + FinishReason *string `thrift:"finish_reason,2,optional" frugal:"2,optional,string" form:"finish_reason" json:"finish_reason,omitempty" query:"finish_reason"` + // token消耗 + Usage *TokenUsage `thrift:"usage,3,optional" frugal:"3,optional,TokenUsage" form:"usage" json:"usage,omitempty" query:"usage"` } func NewExecuteData() *ExecuteData { @@ -2727,11 +2734,14 @@ func (p *ExecuteStreamingResponse) Field255DeepEqual(src *base.BaseResp) bool { } type ExecuteStreamingData struct { - Code *int32 `thrift:"code,1,optional" frugal:"1,optional,i32" form:"code" json:"code,omitempty" query:"code"` - Msg *string `thrift:"msg,2,optional" frugal:"2,optional,string" form:"msg" json:"msg,omitempty" query:"msg"` - Message *Message `thrift:"message,3,optional" frugal:"3,optional,Message" form:"message" json:"message,omitempty" query:"message"` - FinishReason *string `thrift:"finish_reason,4,optional" frugal:"4,optional,string" form:"finish_reason" json:"finish_reason,omitempty" query:"finish_reason"` - Usage *TokenUsage `thrift:"usage,5,optional" frugal:"5,optional,TokenUsage" form:"usage" json:"usage,omitempty" query:"usage"` + Code *int32 `thrift:"code,1,optional" frugal:"1,optional,i32" form:"code" json:"code,omitempty" query:"code"` + Msg *string `thrift:"msg,2,optional" frugal:"2,optional,string" form:"msg" json:"msg,omitempty" query:"msg"` + // 消息 + Message *Message `thrift:"message,3,optional" frugal:"3,optional,Message" form:"message" json:"message,omitempty" query:"message"` + // 结束原因 + FinishReason *string `thrift:"finish_reason,4,optional" frugal:"4,optional,string" form:"finish_reason" json:"finish_reason,omitempty" query:"finish_reason"` + // token消耗 + Usage *TokenUsage `thrift:"usage,5,optional" frugal:"5,optional,TokenUsage" form:"usage" json:"usage,omitempty" query:"usage"` } func NewExecuteStreamingData() *ExecuteStreamingData { @@ -3200,9 +3210,12 @@ func (p *ExecuteStreamingData) Field5DeepEqual(src *TokenUsage) bool { } type PromptQuery struct { + // prompt_key PromptKey *string `thrift:"prompt_key,1,optional" frugal:"1,optional,string" form:"prompt_key" json:"prompt_key,omitempty" query:"prompt_key"` - Version *string `thrift:"version,2,optional" frugal:"2,optional,string" form:"version" json:"version,omitempty" query:"version"` - Label *string `thrift:"label,3,optional" frugal:"3,optional,string" form:"label" json:"label,omitempty" query:"label"` + // prompt版本 + Version *string `thrift:"version,2,optional" frugal:"2,optional,string" form:"version" json:"version,omitempty" query:"version"` + // prompt版本标识（如果version不为空，该字段会被忽略） + Label *string `thrift:"label,3,optional" frugal:"3,optional,string" form:"label" json:"label,omitempty" query:"label"` } func NewPromptQuery() *PromptQuery { @@ -4985,12 +4998,18 @@ func (p *ToolCallConfig) Field1DeepEqual(src *ToolChoiceType) bool { } type Message struct { - Role *Role `thrift:"role,1,optional" frugal:"1,optional,string" form:"role" json:"role,omitempty" query:"role"` - Content *string `thrift:"content,2,optional" frugal:"2,optional,string" form:"content" json:"content,omitempty" query:"content"` - Parts []*ContentPart `thrift:"parts,3,optional" frugal:"3,optional,list" form:"parts" json:"parts,omitempty" query:"parts"` - ReasoningContent *string `thrift:"reasoning_content,4,optional" frugal:"4,optional,string" form:"reasoning_content" json:"reasoning_content,omitempty" query:"reasoning_content"` - ToolCallID *string `thrift:"tool_call_id,5,optional" frugal:"5,optional,string" form:"tool_call_id" json:"tool_call_id,omitempty" query:"tool_call_id"` - ToolCalls []*ToolCall `thrift:"tool_calls,6,optional" frugal:"6,optional,list" form:"tool_calls" json:"tool_calls,omitempty" query:"tool_calls"` + // 角色 + Role *Role `thrift:"role,1,optional" frugal:"1,optional,string" form:"role" json:"role,omitempty" query:"role"` + // 消息内容 + Content *string `thrift:"content,2,optional" frugal:"2,optional,string" form:"content" json:"content,omitempty" query:"content"` + // 多模态内容 + Parts []*ContentPart `thrift:"parts,3,optional" frugal:"3,optional,list" form:"parts" json:"parts,omitempty" query:"parts"` + // 推理思考内容 + ReasoningContent *string `thrift:"reasoning_content,4,optional" frugal:"4,optional,string" form:"reasoning_content" json:"reasoning_content,omitempty" query:"reasoning_content"` + // tool调用ID（role为tool时有效） + ToolCallID *string `thrift:"tool_call_id,5,optional" frugal:"5,optional,string" form:"tool_call_id" json:"tool_call_id,omitempty" query:"tool_call_id"` + // tool调用（role为assistant时有效） + ToolCalls []*ToolCall `thrift:"tool_calls,6,optional" frugal:"6,optional,list" form:"tool_calls" json:"tool_calls,omitempty" query:"tool_calls"` } func NewMessage() *Message { @@ -8233,10 +8252,14 @@ func (p *LLMConfig) Field7DeepEqual(src *bool) bool { } type VariableVal struct { - Key *string `thrift:"key,1,optional" frugal:"1,optional,string" form:"key" json:"key,omitempty" query:"key"` - Value *string `thrift:"value,2,optional" frugal:"2,optional,string" form:"value" json:"value,omitempty" query:"value"` - PlaceholderMessages []*Message `thrift:"placeholder_messages,3,optional" frugal:"3,optional,list" form:"placeholder_messages" json:"placeholder_messages,omitempty" query:"placeholder_messages"` - MultiPartValues []*ContentPart `thrift:"multi_part_values,4,optional" frugal:"4,optional,list" form:"multi_part_values" json:"multi_part_values,omitempty" query:"multi_part_values"` + // 变量key + Key *string `thrift:"key,1,optional" frugal:"1,optional,string" form:"key" json:"key,omitempty" query:"key"` + // 普通变量值（非string类型，如boolean、integer、float、object等，序列化后传入） + Value *string `thrift:"value,2,optional" frugal:"2,optional,string" form:"value" json:"value,omitempty" query:"value"` + // placeholder变量值 + PlaceholderMessages []*Message `thrift:"placeholder_messages,3,optional" frugal:"3,optional,list" form:"placeholder_messages" json:"placeholder_messages,omitempty" query:"placeholder_messages"` + // 多模态变量值 + MultiPartValues []*ContentPart `thrift:"multi_part_values,4,optional" frugal:"4,optional,list" form:"multi_part_values" json:"multi_part_values,omitempty" query:"multi_part_values"` } func NewVariableVal() *VariableVal { @@ -8687,7 +8710,9 @@ func (p *VariableVal) Field4DeepEqual(src []*ContentPart) bool { } type TokenUsage struct { - InputTokens *int32 `thrift:"input_tokens,1,optional" frugal:"1,optional,i32" form:"input_tokens" json:"input_tokens,omitempty" query:"input_tokens"` + // 输入消耗 + InputTokens *int32 `thrift:"input_tokens,1,optional" frugal:"1,optional,i32" form:"input_tokens" json:"input_tokens,omitempty" query:"input_tokens"` + // 输出消耗 OutputTokens *int32 `thrift:"output_tokens,2,optional" frugal:"2,optional,i32" form:"output_tokens" json:"output_tokens,omitempty" query:"output_tokens"` } diff --git a/backend/modules/observability/application/convertor/trace/trace.go b/backend/modules/observability/application/convertor/trace/trace.go index 0ad0f4998..422769a45 100644 --- a/backend/modules/observability/application/convertor/trace/trace.go +++ b/backend/modules/observability/application/convertor/trace/trace.go @@ -34,8 +34,8 @@ func AdvanceInfoDO2TraceDTO(info *loop_span.TraceAdvanceInfo) *traced.Trace { return &traced.Trace{ TraceID: &info.TraceId, Tokens: &traced.TokenCost{ - Input: info.InputCost, - Output: info.OutputCost, + InputToken: info.InputCost, + OutputToken: info.OutputCost, }, } } diff --git a/backend/modules/observability/application/openapi.go b/backend/modules/observability/application/openapi.go index f6c75e9df..8ac414d6b 100644 --- a/backend/modules/observability/application/openapi.go +++ b/backend/modules/observability/application/openapi.go @@ -16,6 +16,7 @@ import ( "github.com/bytedance/gg/gptr" "github.com/bytedance/sonic" + "github.com/coze-dev/coze-loop/backend/kitex_gen/base" coltracepb "go.opentelemetry.io/proto/otlp/collector/trace/v1" "google.golang.org/protobuf/proto" @@ -95,15 +96,15 @@ func (o *OpenAPIApplication) IngestTraces(ctx context.Context, req *openapi.Inge spanMap := o.unpackSpace(ctx, req.Spans) connectorUid := session.UserIDInCtxOrEmpty(ctx) for workspaceId := range spanMap { - // check permission - if err := o.auth.CheckIngestPermission(ctx, workspaceId); err != nil { - return nil, err - } - // check benefit workSpaceIdNum, err := strconv.ParseInt(workspaceId, 10, 64) if err != nil { return nil, errorx.NewByCode(obErrorx.CommercialCommonInvalidParamCodeCode, errorx.WithExtraMsg("invalid workspace_id")) } + // check permission + if err = o.auth.CheckIngestPermission(ctx, workspaceId); err != nil { + return nil, err + } + // check benefit benefitRes, err := o.benefit.CheckTraceBenefit(ctx, &benefit.CheckTraceBenefitParams{ ConnectorUID: connectorUid, SpaceID: workSpaceIdNum, @@ -233,13 +234,13 @@ func (o *OpenAPIApplication) OtelIngestTraces(ctx context.Context, req *openapi. partialFailSpanNumber := 0 partialErrMessage := "" for workspaceId, otelSpans := range spansMap { - if e := o.auth.CheckIngestPermission(ctx, workspaceId); e != nil { - return nil, e - } workSpaceIdNum, e := strconv.ParseInt(workspaceId, 10, 64) if e != nil { return nil, errorx.NewByCode(obErrorx.CommercialCommonInvalidParamCodeCode, errorx.WithExtraMsg("invalid workspace_id")) } + if e = o.auth.CheckIngestPermission(ctx, workspaceId); e != nil { + return nil, e + } connectorUid := session.UserIDInCtxOrEmpty(ctx) benefitRes, e := o.benefit.CheckTraceBenefit(ctx, &benefit.CheckTraceBenefitParams{ ConnectorUID: connectorUid, @@ -463,7 +464,7 @@ func (o *OpenAPIApplication) SearchTraceOApi(ctx context.Context, req *openapi.S errCode := 0 defer func() { if req != nil { - o.metrics.EmitSearchTraceOapi(req.WorkspaceID, req.GetPlatformType(), int64(spansSize), errCode, st, err != nil) + o.metrics.EmitTraceOapi("SearchTraceOApi", req.WorkspaceID, req.GetPlatformType(), "", int64(spansSize), errCode, st, err != nil) } }() @@ -471,12 +472,12 @@ func (o *OpenAPIApplication) SearchTraceOApi(ctx context.Context, req *openapi.S errCode = obErrorx.CommercialCommonInvalidParamCodeCode return nil, err } - req.WorkspaceID = o.workspace.GetQueryWorkSpaceID(ctx, req.GetWorkspaceID()) if err = o.auth.CheckQueryPermission(ctx, strconv.FormatInt(req.GetWorkspaceID(), 10), req.GetPlatformType()); err != nil { errCode = obErrorx.CommonNoPermissionCode return nil, err } - if !o.AllowBySpace(ctx, req.GetWorkspaceID()) { + limitKey := strconv.FormatInt(req.GetWorkspaceID(), 10) + if !o.AllowByKey(ctx, limitKey) { err = errorx.NewByCode(obErrorx.CommonRequestRateLimitCode, errorx.WithExtraMsg("qps limit exceeded")) errCode = obErrorx.CommonRequestRateLimitCode return nil, err @@ -540,14 +541,15 @@ func (o *OpenAPIApplication) buildSearchTraceReq(ctx context.Context, req *opena } ret := &service.SearchTraceOApiReq{ - WorkspaceID: req.WorkspaceID, - Tenants: o.tenant.GetOAPIQueryTenants(ctx, platformType), - TraceID: req.GetTraceID(), - LogID: req.GetLogid(), - StartTime: req.GetStartTime(), - EndTime: req.GetEndTime(), - Limit: req.GetLimit(), - PlatformType: platformType, + WorkspaceID: req.WorkspaceID, + ThirdPartyWorkspaceID: o.workspace.GetThirdPartyQueryWorkSpaceID(ctx, req.WorkspaceID), + Tenants: o.tenant.GetOAPIQueryTenants(ctx, platformType), + TraceID: req.GetTraceID(), + LogID: req.GetLogid(), + StartTime: req.GetStartTime(), + EndTime: req.GetEndTime(), + Limit: req.GetLimit(), + PlatformType: platformType, } if len(ret.Tenants) == 0 { logs.CtxError(ctx, "fail to get platform tenants") @@ -561,22 +563,23 @@ func (o *OpenAPIApplication) ListSpansOApi(ctx context.Context, req *openapi.Lis st := time.Now() spansSize := 0 errCode := 0 + resp := openapi.NewListSpansOApiResponse() defer func() { if req != nil { - o.metrics.EmitListSpansOapi(req.WorkspaceID, req.GetPlatformType(), req.GetSpanListType(), int64(spansSize), errCode, st, err != nil) + o.metrics.EmitTraceOapi("ListSpansOApi", req.WorkspaceID, req.GetPlatformType(), req.GetSpanListType(), int64(spansSize), errCode, st, err != nil) } }() if err = o.validateListSpansOApi(ctx, req); err != nil { errCode = obErrorx.CommercialCommonInvalidParamCodeCode return nil, err } - req.WorkspaceID = o.workspace.GetQueryWorkSpaceID(ctx, req.GetWorkspaceID()) if err = o.auth.CheckQueryPermission(ctx, strconv.FormatInt(req.GetWorkspaceID(), 10), req.GetPlatformType()); err != nil { errCode = obErrorx.CommonNoPermissionCode return nil, err } - if !o.AllowBySpace(ctx, req.GetWorkspaceID()) { + limitKey := strconv.FormatInt(req.GetWorkspaceID(), 10) + if !o.AllowByKey(ctx, limitKey) { err = errorx.NewByCode(obErrorx.CommonRequestRateLimitCode, errorx.WithExtraMsg("qps limit exceeded")) errCode = obErrorx.CommonRequestRateLimitCode return nil, err @@ -595,13 +598,14 @@ func (o *OpenAPIApplication) ListSpansOApi(ctx context.Context, req *openapi.Lis logs.CtxInfo(ctx, "List spans successfully, spans count: %d", len(sResp.Spans)) spansSize = loop_span.SizeofSpans(sResp.Spans) } - return &openapi.ListSpansOApiResponse{ - Data: &openapi.ListSpansOApiData{ - Spans: tconv.SpanListDO2DTO(sResp.Spans, nil, nil, nil), - NextPageToken: sResp.NextPageToken, - HasMore: sResp.HasMore, - }, - }, nil + + resp.Data = &openapi.ListSpansOApiData{ + Spans: tconv.SpanListDO2DTO(sResp.Spans, nil, nil, nil), + NextPageToken: sResp.NextPageToken, + HasMore: sResp.HasMore, + } + resp.BaseResp = base.NewBaseResp() + return resp, nil } func (o *OpenAPIApplication) validateListSpansOApi(ctx context.Context, req *openapi.ListSpansOApiRequest) error { @@ -628,12 +632,13 @@ func (o *OpenAPIApplication) validateListSpansOApi(ctx context.Context, req *ope func (o *OpenAPIApplication) buildListSpansOApiReq(ctx context.Context, req *openapi.ListSpansOApiRequest) (*service.ListSpansOApiReq, error) { ret := &service.ListSpansOApiReq{ - WorkspaceID: req.WorkspaceID, - StartTime: req.GetStartTime(), - EndTime: req.GetEndTime(), - Limit: QueryLimitDefault, - DescByStartTime: len(req.GetOrderBys()) > 0, - PageToken: req.GetPageToken(), + WorkspaceID: req.WorkspaceID, + ThirdPartyWorkspaceID: o.workspace.GetThirdPartyQueryWorkSpaceID(ctx, req.WorkspaceID), + StartTime: req.GetStartTime(), + EndTime: req.GetEndTime(), + Limit: QueryLimitDefault, + DescByStartTime: len(req.GetOrderBys()) > 0, + PageToken: req.GetPageToken(), } if req.PageSize != nil { ret.Limit = *req.PageSize @@ -673,27 +678,27 @@ func (o *OpenAPIApplication) ListTracesOApi(ctx context.Context, req *openapi.Li st := time.Now() errCode := 0 defer func() { - o.metrics.EmitListTracesOapi(req.WorkspaceID, errCode, st, err != nil) + o.metrics.EmitTraceOapi("ListTracesOApi", req.WorkspaceID, "", "", 0, errCode, st, err != nil) }() if err = o.validateListTracesOApiReq(ctx, req); err != nil { errCode = obErrorx.CommercialCommonInvalidParamCodeCode return nil, err } - req.WorkspaceID = o.workspace.GetQueryWorkSpaceID(ctx, req.GetWorkspaceID()) if err = o.auth.CheckQueryPermission(ctx, strconv.FormatInt(req.GetWorkspaceID(), 10), req.GetPlatformType()); err != nil { errCode = obErrorx.CommonNoPermissionCode return nil, err } - if !o.AllowBySpace(ctx, req.GetWorkspaceID()) { + limitKey := strconv.FormatInt(req.GetWorkspaceID(), 10) + if !o.AllowByKey(ctx, limitKey) { err = errorx.NewByCode(obErrorx.CommonRequestRateLimitCode, errorx.WithExtraMsg("qps limit exceeded")) errCode = obErrorx.CommonRequestRateLimitCode return nil, err } logs.CtxInfo(ctx, "ListTracesOApi request: %+v", req) - sReq := o.buildListTracesOApiReq(req) + sReq := o.buildListTracesOApiReq(ctx, req) sResp, err := o.traceService.GetTracesAdvanceInfo(ctx, sReq) if err != nil { errCode = obErrorx.CommonInternalErrorCode @@ -738,10 +743,11 @@ func (o *OpenAPIApplication) validateListTracesOApiReq(ctx context.Context, req return nil } -func (o *OpenAPIApplication) buildListTracesOApiReq(req *openapi.ListTracesOApiRequest) *service.GetTracesAdvanceInfoReq { +func (o *OpenAPIApplication) buildListTracesOApiReq(ctx context.Context, req *openapi.ListTracesOApiRequest) *service.GetTracesAdvanceInfoReq { ret := &service.GetTracesAdvanceInfoReq{ - WorkspaceID: req.GetWorkspaceID(), - Traces: make([]*service.TraceQueryParam, len(req.GetTraceIds())), + WorkspaceID: req.GetWorkspaceID(), + ThirdPartyWorkspaceID: o.workspace.GetThirdPartyQueryWorkSpaceID(ctx, req.WorkspaceID), + Traces: make([]*service.TraceQueryParam, len(req.GetTraceIds())), } for i, id := range req.GetTraceIds() { ret.Traces[i] = &service.TraceQueryParam{ @@ -762,13 +768,13 @@ func (o *OpenAPIApplication) Send(ctx context.Context, event *entity.AnnotationE return o.traceService.Send(ctx, event) } -func (p *OpenAPIApplication) AllowBySpace(ctx context.Context, workspaceID int64) bool { - maxQPS, err := p.traceConfig.GetQueryMaxQPSBySpace(ctx, workspaceID) +func (p *OpenAPIApplication) AllowByKey(ctx context.Context, key string) bool { + maxQPS, err := p.traceConfig.GetQueryMaxQPS(ctx, key) if err != nil { - logs.CtxError(ctx, "get query max qps failed, err=%v, space_id=%d", err, workspaceID) + logs.CtxError(ctx, "get query max qps failed, err=%v, key=%s", err, key) return true } - result, err := p.rateLimiter.AllowN(ctx, fmt.Sprintf("query_trace:qps:space_id:%d", workspaceID), 1, + result, err := p.rateLimiter.AllowN(ctx, key, 1, limiter.WithLimit(&limiter.Limit{ Rate: maxQPS, Burst: maxQPS, diff --git a/backend/modules/observability/application/openapi_test.go b/backend/modules/observability/application/openapi_test.go index 2b0599137..7aed27b68 100644 --- a/backend/modules/observability/application/openapi_test.go +++ b/backend/modules/observability/application/openapi_test.go @@ -84,12 +84,12 @@ func TestOpenAPIApplication_IngestTraces(t *testing.T) { tenantMock := tenantmocks.NewMockITenantProvider(ctrl) tenantMock.EXPECT().GetIngestTenant(gomock.Any(), gomock.Any()).Return("t") workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("1") + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("1").AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), gomock.Any()).Return(100, nil).AnyTimes() + traceConfigMock.EXPECT().GetQueryMaxQPS(gomock.Any(), gomock.Any()).Return(100, nil).AnyTimes() traceConfigMock.EXPECT().GetTraceIngestTenantProducerCfg(gomock.Any()).Return(nil, nil).AnyTimes() return fields{ traceService: traceServiceMock, @@ -153,7 +153,17 @@ func TestOpenAPIApplication_IngestTraces(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() fields := tt.fieldsGetter(ctrl) - o, err := NewOpenAPIApplication(fields.traceService, fields.auth, fields.benefit, fields.tenant, fields.workspace, fields.rateLimiter, fields.traceConfig, fields.metrics) + o := &OpenAPIApplication{ + traceService: fields.traceService, + auth: fields.auth, + benefit: fields.benefit, + tenant: fields.tenant, + workspace: fields.workspace, + rateLimiter: fields.rateLimiter.NewRateLimiter(), + traceConfig: fields.traceConfig, + metrics: fields.metrics, + } + err := error(nil) assert.NoError(t, err) got, err := o.IngestTraces(tt.args.ctx, tt.args.req) assert.Equal(t, tt.wantErr, err != nil) @@ -196,6 +206,18 @@ func TestOpenAPIApplication_CreateAnnotation(t *testing.T) { authMock := rpcmocks.NewMockIAuthProvider(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -231,6 +253,18 @@ func TestOpenAPIApplication_CreateAnnotation(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -264,7 +298,17 @@ func TestOpenAPIApplication_CreateAnnotation(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() fields := tt.fieldsGetter(ctrl) - o, err := NewOpenAPIApplication(fields.traceService, fields.auth, fields.benefit, fields.tenant, fields.workspace, fields.rateLimiter, fields.traceConfig, fields.metrics) + o := &OpenAPIApplication{ + traceService: fields.traceService, + auth: fields.auth, + benefit: fields.benefit, + tenant: fields.tenant, + workspace: fields.workspace, + rateLimiter: fields.rateLimiter.NewRateLimiter(), + traceConfig: fields.traceConfig, + metrics: fields.metrics, + } + err := error(nil) assert.NoError(t, err) got, err := o.CreateAnnotation(tt.args.ctx, tt.args.req) assert.Equal(t, tt.wantErr, err != nil) @@ -307,6 +351,18 @@ func TestOpenAPIApplication_DeleteAnnotation(t *testing.T) { authMock := rpcmocks.NewMockIAuthProvider(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -338,7 +394,17 @@ func TestOpenAPIApplication_DeleteAnnotation(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() fields := tt.fieldsGetter(ctrl) - o, err := NewOpenAPIApplication(fields.traceService, fields.auth, fields.benefit, fields.tenant, fields.workspace, fields.rateLimiter, fields.traceConfig, fields.metrics) + o := &OpenAPIApplication{ + traceService: fields.traceService, + auth: fields.auth, + benefit: fields.benefit, + tenant: fields.tenant, + workspace: fields.workspace, + rateLimiter: fields.rateLimiter.NewRateLimiter(), + traceConfig: fields.traceConfig, + metrics: fields.metrics, + } + err := error(nil) assert.NoError(t, err) got, err := o.DeleteAnnotation(tt.args.ctx, tt.args.req) assert.Equal(t, tt.wantErr, err != nil) @@ -377,6 +443,18 @@ func TestOpenAPIApplication_Send(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -404,7 +482,17 @@ func TestOpenAPIApplication_Send(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() fields := tt.fieldsGetter(ctrl) - o, err := NewOpenAPIApplication(fields.traceService, fields.auth, fields.benefit, fields.tenant, fields.workspace, fields.rateLimiter, fields.traceConfig, fields.metrics) + o := &OpenAPIApplication{ + traceService: fields.traceService, + auth: fields.auth, + benefit: fields.benefit, + tenant: fields.tenant, + workspace: fields.workspace, + rateLimiter: fields.rateLimiter.NewRateLimiter(), + traceConfig: fields.traceConfig, + metrics: fields.metrics, + } + err := error(nil) assert.NoError(t, err) err = o.Send(tt.args.ctx, tt.args.event) assert.Equal(t, tt.wantErr, err != nil) @@ -451,6 +539,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { tenantMock := tenantmocks.NewMockITenantProvider(ctrl) tenantMock.EXPECT().GetIngestTenant(gomock.Any(), gomock.Any()).Return("test-tenant") workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -498,6 +598,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { tenantMock := tenantmocks.NewMockITenantProvider(ctrl) tenantMock.EXPECT().GetIngestTenant(gomock.Any(), gomock.Any()).Return("test-tenant") workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -545,6 +657,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { tenantMock := tenantmocks.NewMockITenantProvider(ctrl) tenantMock.EXPECT().GetIngestTenant(gomock.Any(), gomock.Any()).Return("test-tenant") workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -588,6 +712,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { tenantMock := tenantmocks.NewMockITenantProvider(ctrl) tenantMock.EXPECT().GetIngestTenant(gomock.Any(), gomock.Any()).Return("test-tenant") workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -626,6 +762,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -656,6 +804,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -691,6 +851,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -726,6 +898,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -761,6 +945,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -796,6 +992,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -828,10 +1036,22 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { fieldsGetter: func(ctrl *gomock.Controller) fields { traceServiceMock := servicemocks.NewMockITraceService(ctrl) authMock := rpcmocks.NewMockIAuthProvider(ctrl) - authMock.EXPECT().CheckIngestPermission(gomock.Any(), "invalid").Return(nil) + authMock.EXPECT().CheckIngestPermission(gomock.Any(), "invalid").Return(nil).AnyTimes() benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -868,6 +1088,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -910,6 +1142,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { }, nil) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -952,6 +1196,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { }, nil) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -996,6 +1252,18 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { tenantMock := tenantmocks.NewMockITenantProvider(ctrl) tenantMock.EXPECT().GetIngestTenant(gomock.Any(), gomock.Any()).Return("test-tenant") workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -1029,7 +1297,17 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() fields := tt.fieldsGetter(ctrl) - o, err := NewOpenAPIApplication(fields.traceService, fields.auth, fields.benefit, fields.tenant, fields.workspace, fields.rateLimiter, fields.traceConfig, fields.metrics) + o := &OpenAPIApplication{ + traceService: fields.traceService, + auth: fields.auth, + benefit: fields.benefit, + tenant: fields.tenant, + workspace: fields.workspace, + rateLimiter: fields.rateLimiter.NewRateLimiter(), + traceConfig: fields.traceConfig, + metrics: fields.metrics, + } + err := error(nil) assert.NoError(t, err) got, err := o.OtelIngestTraces(tt.args.ctx, tt.args.req) assert.Equal(t, tt.wantErr, err != nil) @@ -1045,10 +1323,6 @@ func TestOpenAPIApplication_OtelIngestTraces(t *testing.T) { } } -// Test helper functions - -// Test helper functions - // createValidJSONTraceData creates valid JSON format trace data for testing func createValidJSONTraceData() []byte { req := &otel.ExportTraceServiceRequest{ @@ -1203,13 +1477,25 @@ func TestOpenAPIApplication_ListSpansOApi(t *testing.T) { tenantMock := tenantmocks.NewMockITenantProvider(ctrl) tenantMock.EXPECT().GetOAPIQueryTenants(gomock.Any(), gomock.Any()).Return([]string{"tenant1"}) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterFactoryMock := limitermocks.NewMockIRateLimiter(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitListSpansOapi( + metricsMock.EXPECT().EmitTraceOapi( + "ListSpansOApi", int64(123), gomock.Any(), gomock.Any(), @@ -1218,7 +1504,7 @@ func TestOpenAPIApplication_ListSpansOApi(t *testing.T) { gomock.Any(), gomock.Any(), ).AnyTimes() - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) + traceConfigMock.EXPECT().GetQueryMaxQPS(gomock.Any(), "123").Return(100, nil) rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&limiter.Result{Allowed: true}, nil) return fields{ traceService: traceServiceMock, @@ -1249,36 +1535,6 @@ func TestOpenAPIApplication_ListSpansOApi(t *testing.T) { }, wantErr: false, }, - { - name: "request is nil", - fieldsGetter: func(ctrl *gomock.Controller) fields { - traceServiceMock := servicemocks.NewMockITraceService(ctrl) - authMock := rpcmocks.NewMockIAuthProvider(ctrl) - benefitMock := benefitmocks.NewMockIBenefitService(ctrl) - tenantMock := tenantmocks.NewMockITenantProvider(ctrl) - workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) - rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() - traceConfigMock := configmocks.NewMockITraceConfig(ctrl) - metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - return fields{ - traceService: traceServiceMock, - auth: authMock, - benefit: benefitMock, - tenant: tenantMock, - workspace: workspaceMock, - rateLimiter: rateLimiterMock, - traceConfig: traceConfigMock, - metrics: metricsMock, - } - }, - args: args{ - ctx: context.Background(), - req: nil, - }, - want: nil, - wantErr: true, - }, { name: "page size exceeds limit", fieldsGetter: func(ctrl *gomock.Controller) fields { @@ -1287,11 +1543,24 @@ func TestOpenAPIApplication_ListSpansOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitListSpansOapi( + metricsMock.EXPECT().EmitTraceOapi( + "ListSpansOApi", int64(123), gomock.Any(), gomock.Any(), @@ -1332,12 +1601,24 @@ func TestOpenAPIApplication_ListSpansOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitListSpansOapi( + metricsMock.EXPECT().EmitTraceOapi( + "ListSpansOApi", int64(123), gomock.Any(), gomock.Any(), @@ -1381,11 +1662,23 @@ func TestOpenAPIApplication_ListSpansOApi(t *testing.T) { rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) + traceConfigMock.EXPECT().GetQueryMaxQPS(gomock.Any(), "123").Return(100, nil) rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&limiter.Result{Allowed: false}, nil) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) - metricsMock.EXPECT().EmitListSpansOapi( + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() + metricsMock.EXPECT().EmitTraceOapi( + "ListSpansOApi", int64(123), gomock.Any(), gomock.Any(), @@ -1431,11 +1724,23 @@ func TestOpenAPIApplication_ListSpansOApi(t *testing.T) { rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) + traceConfigMock.EXPECT().GetQueryMaxQPS(gomock.Any(), "123").Return(100, nil) rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&limiter.Result{Allowed: true}, nil) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) - metricsMock.EXPECT().EmitListSpansOapi( + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() + metricsMock.EXPECT().EmitTraceOapi( + "ListSpansOApi", int64(123), gomock.Any(), gomock.Any(), @@ -1472,7 +1777,17 @@ func TestOpenAPIApplication_ListSpansOApi(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() fields := tt.fieldsGetter(ctrl) - o, err := NewOpenAPIApplication(fields.traceService, fields.auth, fields.benefit, fields.tenant, fields.workspace, fields.rateLimiter, fields.traceConfig, fields.metrics) + o := &OpenAPIApplication{ + traceService: fields.traceService, + auth: fields.auth, + benefit: fields.benefit, + tenant: fields.tenant, + workspace: fields.workspace, + rateLimiter: fields.rateLimiter.NewRateLimiter(), + traceConfig: fields.traceConfig, + metrics: fields.metrics, + } + err := error(nil) assert.NoError(t, err) got, err := o.ListSpansOApi(tt.args.ctx, tt.args.req) assert.Equal(t, tt.wantErr, err != nil) @@ -1526,13 +1841,26 @@ func TestOpenAPIApplication_SearchTraceOApi(t *testing.T) { rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) + traceConfigMock.EXPECT().GetQueryMaxQPS(gomock.Any(), "123").Return(100, nil) rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&limiter.Result{Allowed: true}, nil) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) - metricsMock.EXPECT().EmitSearchTraceOapi( + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() + metricsMock.EXPECT().EmitTraceOapi( + "SearchTraceOApi", int64(123), gomock.Any(), + "", gomock.Any(), gomock.Any(), gomock.Any(), @@ -1583,13 +1911,26 @@ func TestOpenAPIApplication_SearchTraceOApi(t *testing.T) { rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) + traceConfigMock.EXPECT().GetQueryMaxQPS(gomock.Any(), "123").Return(100, nil) rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&limiter.Result{Allowed: true}, nil) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) - metricsMock.EXPECT().EmitSearchTraceOapi( + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() + metricsMock.EXPECT().EmitTraceOapi( + "SearchTraceOApi", int64(123), gomock.Any(), + "", gomock.Any(), gomock.Any(), gomock.Any(), @@ -1631,6 +1972,18 @@ func TestOpenAPIApplication_SearchTraceOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -1661,13 +2014,27 @@ func TestOpenAPIApplication_SearchTraceOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitSearchTraceOapi( + metricsMock.EXPECT().EmitTraceOapi( + "SearchTraceOApi", int64(123), gomock.Any(), + "", gomock.Any(), gomock.Any(), gomock.Any(), @@ -1704,13 +2071,27 @@ func TestOpenAPIApplication_SearchTraceOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitSearchTraceOapi( + metricsMock.EXPECT().EmitTraceOapi( + "SearchTraceOApi", int64(123), gomock.Any(), + "", gomock.Any(), gomock.Any(), gomock.Any(), @@ -1749,14 +2130,27 @@ func TestOpenAPIApplication_SearchTraceOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitSearchTraceOapi( + metricsMock.EXPECT().EmitTraceOapi( + "SearchTraceOApi", int64(123), gomock.Any(), + "", gomock.Any(), gomock.Any(), gomock.Any(), @@ -1801,13 +2195,26 @@ func TestOpenAPIApplication_SearchTraceOApi(t *testing.T) { rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) + traceConfigMock.EXPECT().GetQueryMaxQPS(gomock.Any(), "123").Return(100, nil) rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&limiter.Result{Allowed: true}, nil) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) - metricsMock.EXPECT().EmitSearchTraceOapi( + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() + metricsMock.EXPECT().EmitTraceOapi( + "SearchTraceOApi", int64(123), gomock.Any(), + "", gomock.Any(), gomock.Any(), gomock.Any(), @@ -1843,7 +2250,17 @@ func TestOpenAPIApplication_SearchTraceOApi(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() fields := tt.fieldsGetter(ctrl) - o, err := NewOpenAPIApplication(fields.traceService, fields.auth, fields.benefit, fields.tenant, fields.workspace, fields.rateLimiter, fields.traceConfig, fields.metrics) + o := &OpenAPIApplication{ + traceService: fields.traceService, + auth: fields.auth, + benefit: fields.benefit, + tenant: fields.tenant, + workspace: fields.workspace, + rateLimiter: fields.rateLimiter.NewRateLimiter(), + traceConfig: fields.traceConfig, + metrics: fields.metrics, + } + err := error(nil) assert.NoError(t, err) got, err := o.SearchTraceOApi(tt.args.ctx, tt.args.req) assert.Equal(t, tt.wantErr, err != nil) @@ -1903,19 +2320,34 @@ func TestOpenAPIApplication_ListTracesOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterFactoryMock := limitermocks.NewMockIRateLimiter(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitListTracesOapi( + metricsMock.EXPECT().EmitTraceOapi( + "ListTracesOApi", int64(123), + "", + "", + int64(0), gomock.Any(), gomock.Any(), gomock.Any(), ).AnyTimes() - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) + traceConfigMock.EXPECT().GetQueryMaxQPS(gomock.Any(), "123").Return(100, nil) rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&limiter.Result{Allowed: true}, nil) return fields{ traceService: traceServiceMock, @@ -1959,11 +2391,27 @@ func TestOpenAPIApplication_ListTracesOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitListTracesOapi( + metricsMock.EXPECT().EmitTraceOapi( + "ListTracesOApi", + int64(0), + "", + "", int64(0), gomock.Any(), gomock.Any(), @@ -2000,12 +2448,28 @@ func TestOpenAPIApplication_ListTracesOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitListTracesOapi( + metricsMock.EXPECT().EmitTraceOapi( + "ListTracesOApi", int64(123), + "", + "", + int64(0), gomock.Any(), gomock.Any(), gomock.Any(), @@ -2041,12 +2505,28 @@ func TestOpenAPIApplication_ListTracesOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitListTracesOapi( + metricsMock.EXPECT().EmitTraceOapi( + "ListTracesOApi", int64(123), + "", + "", + int64(0), gomock.Any(), gomock.Any(), gomock.Any(), @@ -2083,13 +2563,28 @@ func TestOpenAPIApplication_ListTracesOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitListTracesOapi( + metricsMock.EXPECT().EmitTraceOapi( + "ListTracesOApi", int64(123), + "", + "", + int64(0), gomock.Any(), gomock.Any(), gomock.Any(), @@ -2126,19 +2621,34 @@ func TestOpenAPIApplication_ListTracesOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterFactoryMock := limitermocks.NewMockIRateLimiter(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitListTracesOapi( + metricsMock.EXPECT().EmitTraceOapi( + "ListTracesOApi", int64(123), + "", + "", + int64(0), gomock.Any(), gomock.Any(), gomock.Any(), ).AnyTimes() - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) + traceConfigMock.EXPECT().GetQueryMaxQPS(gomock.Any(), "123").Return(100, nil) rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&limiter.Result{Allowed: false}, nil) return fields{ traceService: traceServiceMock, @@ -2173,19 +2683,34 @@ func TestOpenAPIApplication_ListTracesOApi(t *testing.T) { benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetQueryWorkSpaceID(gomock.Any(), int64(123)).Return(int64(123)) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).DoAndReturn(func(ctx context.Context, spans []*span.InputSpan) string { + if len(spans) > 0 { + switch spans[0].SpanID { + case "span1": + case "span2": + case "span3": + return "workspace2" + } + } + return "" + }).AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterFactoryMock := limitermocks.NewMockIRateLimiter(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - metricsMock.EXPECT().EmitListTracesOapi( + metricsMock.EXPECT().EmitTraceOapi( + "ListTracesOApi", int64(123), + "", + "", + int64(0), gomock.Any(), gomock.Any(), gomock.Any(), ).AnyTimes() - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) + traceConfigMock.EXPECT().GetQueryMaxQPS(gomock.Any(), "123").Return(100, nil) rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(&limiter.Result{Allowed: true}, nil) return fields{ traceService: traceServiceMock, @@ -2217,7 +2742,17 @@ func TestOpenAPIApplication_ListTracesOApi(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() fields := tt.fieldsGetter(ctrl) - o, err := NewOpenAPIApplication(fields.traceService, fields.auth, fields.benefit, fields.tenant, fields.workspace, fields.rateLimiter, fields.traceConfig, fields.metrics) + o := &OpenAPIApplication{ + traceService: fields.traceService, + auth: fields.auth, + benefit: fields.benefit, + tenant: fields.tenant, + workspace: fields.workspace, + rateLimiter: fields.rateLimiter.NewRateLimiter(), + traceConfig: fields.traceConfig, + metrics: fields.metrics, + } + err := error(nil) assert.NoError(t, err) got, err := o.ListTracesOApi(tt.args.ctx, tt.args.req) assert.Equal(t, tt.wantErr, err != nil) @@ -2233,7 +2768,55 @@ func TestOpenAPIApplication_ListTracesOApi(t *testing.T) { } // TestOpenAPIApplication_unpackSpace tests the unpackSpace method -func TestOpenAPIApplication_unpackSpace(t *testing.T) { + +// TestOpenAPIApplication_AllowByKey tests the AllowByKey method + +// TestNewOpenAPIApplication 测试构造函数 +func TestNewOpenAPIApplication(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + traceServiceMock := servicemocks.NewMockITraceService(ctrl) + authMock := rpcmocks.NewMockIAuthProvider(ctrl) + benefitMock := benefitmocks.NewMockIBenefitService(ctrl) + tenantMock := tenantmocks.NewMockITenantProvider(ctrl) + workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + rateLimiterFactoryMock := limitermocks.NewMockIRateLimiterFactory(ctrl) + rateLimiterMock := limitermocks.NewMockIRateLimiter(ctrl) + traceConfigMock := configmocks.NewMockITraceConfig(ctrl) + metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) + + rateLimiterFactoryMock.EXPECT().NewRateLimiter().Return(rateLimiterMock) + + app, err := NewOpenAPIApplication( + traceServiceMock, + authMock, + benefitMock, + tenantMock, + workspaceMock, + rateLimiterFactoryMock, + traceConfigMock, + metricsMock, + ) + + assert.NoError(t, err) + assert.NotNil(t, app) + + // 验证返回的实例类型 + openAPIApp, ok := app.(*OpenAPIApplication) + assert.True(t, ok) + assert.NotNil(t, openAPIApp.traceService) + assert.NotNil(t, openAPIApp.auth) + assert.NotNil(t, openAPIApp.benefit) + assert.NotNil(t, openAPIApp.tenant) + assert.NotNil(t, openAPIApp.workspace) + assert.NotNil(t, openAPIApp.rateLimiter) + assert.NotNil(t, openAPIApp.traceConfig) + assert.NotNil(t, openAPIApp.metrics) +} + +// 补充IngestTraces的边界测试场景 +func TestOpenAPIApplication_IngestTraces_AdditionalScenarios(t *testing.T) { type fields struct { traceService service.ITraceService auth rpc.IAuthProvider @@ -2245,23 +2828,26 @@ func TestOpenAPIApplication_unpackSpace(t *testing.T) { metrics metrics.ITraceMetrics } type args struct { - ctx context.Context - spans []*span.InputSpan + ctx context.Context + req *openapi.IngestTracesRequest } tests := []struct { name string fieldsGetter func(ctrl *gomock.Controller) fields args args - want map[string][]*span.InputSpan + want *openapi.IngestTracesResponse + wantErr bool }{ { - name: "nil spans", + name: "permission check fails", fieldsGetter: func(ctrl *gomock.Controller) fields { traceServiceMock := servicemocks.NewMockITraceService(ctrl) authMock := rpcmocks.NewMockIAuthProvider(ctrl) + authMock.EXPECT().CheckIngestPermission(gomock.Any(), gomock.Any()).Return(assert.AnError) benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("1").AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -2278,20 +2864,33 @@ func TestOpenAPIApplication_unpackSpace(t *testing.T) { } }, args: args{ - ctx: context.Background(), - spans: nil, + ctx: context.Background(), + req: &openapi.IngestTracesRequest{ + Spans: []*span.InputSpan{ + { + WorkspaceID: "1", + }, + }, + }, }, - want: nil, + want: nil, + wantErr: true, }, { - name: "empty workspace ID skipped", + name: "benefit check fails - insufficient capacity", fieldsGetter: func(ctrl *gomock.Controller) fields { traceServiceMock := servicemocks.NewMockITraceService(ctrl) authMock := rpcmocks.NewMockIAuthProvider(ctrl) + authMock.EXPECT().CheckIngestPermission(gomock.Any(), gomock.Any()).Return(nil) benefitMock := benefitmocks.NewMockIBenefitService(ctrl) + benefitMock.EXPECT().CheckTraceBenefit(gomock.Any(), gomock.Any()).Return(&benefit.CheckTraceBenefitResult{ + AccountAvailable: true, + IsEnough: false, + StorageDuration: 3, + }, nil) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("") + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("1").AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -2309,22 +2908,32 @@ func TestOpenAPIApplication_unpackSpace(t *testing.T) { }, args: args{ ctx: context.Background(), - spans: []*span.InputSpan{ - {SpanID: "span1"}, + req: &openapi.IngestTracesRequest{ + Spans: []*span.InputSpan{ + { + WorkspaceID: "1", + }, + }, }, }, - want: map[string][]*span.InputSpan{}, + want: nil, + wantErr: true, }, { - name: "spans grouped by workspace ID", + name: "benefit check fails - account not available", fieldsGetter: func(ctrl *gomock.Controller) fields { traceServiceMock := servicemocks.NewMockITraceService(ctrl) authMock := rpcmocks.NewMockIAuthProvider(ctrl) + authMock.EXPECT().CheckIngestPermission(gomock.Any(), gomock.Any()).Return(nil) benefitMock := benefitmocks.NewMockIBenefitService(ctrl) + benefitMock.EXPECT().CheckTraceBenefit(gomock.Any(), gomock.Any()).Return(&benefit.CheckTraceBenefitResult{ + AccountAvailable: false, + IsEnough: true, + StorageDuration: 3, + }, nil) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) - workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("workspace1").Times(2) - workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("workspace2").Times(1) + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("1").AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) @@ -2342,48 +2951,110 @@ func TestOpenAPIApplication_unpackSpace(t *testing.T) { }, args: args{ ctx: context.Background(), - spans: []*span.InputSpan{ - {SpanID: "span1"}, - {SpanID: "span2"}, - {SpanID: "span3"}, + req: &openapi.IngestTracesRequest{ + Spans: []*span.InputSpan{ + { + WorkspaceID: "1", + }, + }, }, }, - want: map[string][]*span.InputSpan{ - "workspace1": { - {SpanID: "span1", WorkspaceID: "workspace1"}, - {SpanID: "span2", WorkspaceID: "workspace1"}, - }, - "workspace2": { - {SpanID: "span3", WorkspaceID: "workspace2"}, + want: nil, + wantErr: true, + }, + { + name: "invalid workspace id format", + fieldsGetter: func(ctrl *gomock.Controller) fields { + traceServiceMock := servicemocks.NewMockITraceService(ctrl) + authMock := rpcmocks.NewMockIAuthProvider(ctrl) + authMock.EXPECT().CheckIngestPermission(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() + benefitMock := benefitmocks.NewMockIBenefitService(ctrl) + tenantMock := tenantmocks.NewMockITenantProvider(ctrl) + workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("invalid").AnyTimes() + rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) + rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() + traceConfigMock := configmocks.NewMockITraceConfig(ctrl) + metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) + return fields{ + traceService: traceServiceMock, + auth: authMock, + benefit: benefitMock, + tenant: tenantMock, + workspace: workspaceMock, + rateLimiter: rateLimiterMock, + traceConfig: traceConfigMock, + metrics: metricsMock, + } + }, + args: args{ + ctx: context.Background(), + req: &openapi.IngestTracesRequest{ + Spans: []*span.InputSpan{ + { + WorkspaceID: "1", + }, + }, }, }, + want: nil, + wantErr: true, + }, + { + name: "nil request", + fieldsGetter: func(ctrl *gomock.Controller) fields { + traceServiceMock := servicemocks.NewMockITraceService(ctrl) + authMock := rpcmocks.NewMockIAuthProvider(ctrl) + benefitMock := benefitmocks.NewMockIBenefitService(ctrl) + tenantMock := tenantmocks.NewMockITenantProvider(ctrl) + workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) + rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() + traceConfigMock := configmocks.NewMockITraceConfig(ctrl) + metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) + return fields{ + traceService: traceServiceMock, + auth: authMock, + benefit: benefitMock, + tenant: tenantMock, + workspace: workspaceMock, + rateLimiter: rateLimiterMock, + traceConfig: traceConfigMock, + metrics: metricsMock, + } + }, + args: args{ + ctx: context.Background(), + req: nil, + }, + want: nil, + wantErr: true, }, } - for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() fields := tt.fieldsGetter(ctrl) - o, err := NewOpenAPIApplication(fields.traceService, fields.auth, fields.benefit, fields.tenant, fields.workspace, fields.rateLimiter, fields.traceConfig, fields.metrics) - assert.NoError(t, err) - got := o.(*OpenAPIApplication).unpackSpace(tt.args.ctx, tt.args.spans) - if tt.want == nil { - assert.Nil(t, got) - } else { - assert.Equal(t, len(tt.want), len(got)) - for workspaceID, expectedSpans := range tt.want { - actualSpans, exists := got[workspaceID] - assert.True(t, exists) - assert.Equal(t, len(expectedSpans), len(actualSpans)) - } + o := &OpenAPIApplication{ + traceService: fields.traceService, + auth: fields.auth, + benefit: fields.benefit, + tenant: fields.tenant, + workspace: fields.workspace, + rateLimiter: fields.rateLimiter.NewRateLimiter(), + traceConfig: fields.traceConfig, + metrics: fields.metrics, } + got, err := o.IngestTraces(tt.args.ctx, tt.args.req) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) }) } } -// TestOpenAPIApplication_AllowBySpace tests the AllowBySpace method -func TestOpenAPIApplication_AllowBySpace(t *testing.T) { +// 补充CreateAnnotation的更多测试场景 +func TestOpenAPIApplication_CreateAnnotation_AdditionalScenarios(t *testing.T) { type fields struct { traceService service.ITraceService auth rpc.IAuthProvider @@ -2395,30 +3066,75 @@ func TestOpenAPIApplication_AllowBySpace(t *testing.T) { metrics metrics.ITraceMetrics } type args struct { - ctx context.Context - workspaceID int64 + ctx context.Context + req *openapi.CreateAnnotationRequest } tests := []struct { name string fieldsGetter func(ctrl *gomock.Controller) fields args args - want bool + want *openapi.CreateAnnotationResponse + wantErr bool }{ { - name: "rate limit allowed", + name: "create annotation with double value type", fieldsGetter: func(ctrl *gomock.Controller) fields { traceServiceMock := servicemocks.NewMockITraceService(ctrl) + traceServiceMock.EXPECT().CreateAnnotation(gomock.Any(), gomock.Any()).Return(nil) + benefitMock := benefitmocks.NewMockIBenefitService(ctrl) + benefitMock.EXPECT().CheckTraceBenefit(gomock.Any(), gomock.Any()).Return(&benefit.CheckTraceBenefitResult{ + StorageDuration: 3, + }, nil) authMock := rpcmocks.NewMockIAuthProvider(ctrl) + tenantMock := tenantmocks.NewMockITenantProvider(ctrl) + workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("").AnyTimes() + rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) + rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() + traceConfigMock := configmocks.NewMockITraceConfig(ctrl) + metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) + return fields{ + traceService: traceServiceMock, + auth: authMock, + benefit: benefitMock, + tenant: tenantMock, + workspace: workspaceMock, + rateLimiter: rateLimiterMock, + traceConfig: traceConfigMock, + metrics: metricsMock, + } + }, + args: args{ + ctx: context.Background(), + req: &openapi.CreateAnnotationRequest{ + WorkspaceID: 1, + AnnotationValueType: ptr.Of(annotation.ValueType(loop_span.AnnotationValueTypeDouble)), + AnnotationValue: "3.14", + Base: &base.Base{Caller: "test"}, + }, + }, + want: openapi.NewCreateAnnotationResponse(), + wantErr: false, + }, + { + name: "create annotation with bool value type", + fieldsGetter: func(ctrl *gomock.Controller) fields { + traceServiceMock := servicemocks.NewMockITraceService(ctrl) + traceServiceMock.EXPECT().CreateAnnotation(gomock.Any(), gomock.Any()).Return(nil) benefitMock := benefitmocks.NewMockIBenefitService(ctrl) + benefitMock.EXPECT().CheckTraceBenefit(gomock.Any(), gomock.Any()).Return(&benefit.CheckTraceBenefitResult{ + StorageDuration: 3, + }, nil) + authMock := rpcmocks.NewMockIAuthProvider(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("").AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) - rateLimiterFactoryMock := limitermocks.NewMockIRateLimiter(ctrl) - rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() + rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) - rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), "query_trace:qps:space_id:123", 1, gomock.Any()).Return(&limiter.Result{Allowed: true}, nil) return fields{ traceService: traceServiceMock, auth: authMock, @@ -2431,26 +3147,31 @@ func TestOpenAPIApplication_AllowBySpace(t *testing.T) { } }, args: args{ - ctx: context.Background(), - workspaceID: 123, + ctx: context.Background(), + req: &openapi.CreateAnnotationRequest{ + WorkspaceID: 1, + AnnotationValueType: ptr.Of(annotation.ValueType(loop_span.AnnotationValueTypeBool)), + AnnotationValue: "true", + Base: &base.Base{Caller: "test"}, + }, }, - want: true, + want: openapi.NewCreateAnnotationResponse(), + wantErr: false, }, { - name: "rate limit exceeded", + name: "create annotation with invalid double value", fieldsGetter: func(ctrl *gomock.Controller) fields { traceServiceMock := servicemocks.NewMockITraceService(ctrl) authMock := rpcmocks.NewMockIAuthProvider(ctrl) benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("").AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) - rateLimiterFactoryMock := limitermocks.NewMockIRateLimiter(ctrl) - rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() + rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) - rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), "query_trace:qps:space_id:123", 1, gomock.Any()).Return(&limiter.Result{Allowed: false}, nil) return fields{ traceService: traceServiceMock, auth: authMock, @@ -2463,25 +3184,31 @@ func TestOpenAPIApplication_AllowBySpace(t *testing.T) { } }, args: args{ - ctx: context.Background(), - workspaceID: 123, + ctx: context.Background(), + req: &openapi.CreateAnnotationRequest{ + WorkspaceID: 1, + AnnotationValueType: ptr.Of(annotation.ValueType(loop_span.AnnotationValueTypeDouble)), + AnnotationValue: "invalid", + Base: &base.Base{Caller: "test"}, + }, }, - want: false, + want: nil, + wantErr: true, }, { - name: "config error returns true", + name: "create annotation with invalid bool value", fieldsGetter: func(ctrl *gomock.Controller) fields { traceServiceMock := servicemocks.NewMockITraceService(ctrl) authMock := rpcmocks.NewMockIAuthProvider(ctrl) benefitMock := benefitmocks.NewMockIBenefitService(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("").AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) - rateLimiterFactoryMock := limitermocks.NewMockIRateLimiter(ctrl) - rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() + rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(0, assert.AnError) return fields{ traceService: traceServiceMock, auth: authMock, @@ -2494,26 +3221,32 @@ func TestOpenAPIApplication_AllowBySpace(t *testing.T) { } }, args: args{ - ctx: context.Background(), - workspaceID: 123, + ctx: context.Background(), + req: &openapi.CreateAnnotationRequest{ + WorkspaceID: 1, + AnnotationValueType: ptr.Of(annotation.ValueType(loop_span.AnnotationValueTypeBool)), + AnnotationValue: "invalid", + Base: &base.Base{Caller: "test"}, + }, }, - want: true, + want: nil, + wantErr: true, }, { - name: "rate limiter error returns true", + name: "benefit check fails", fieldsGetter: func(ctrl *gomock.Controller) fields { traceServiceMock := servicemocks.NewMockITraceService(ctrl) authMock := rpcmocks.NewMockIAuthProvider(ctrl) benefitMock := benefitmocks.NewMockIBenefitService(ctrl) + benefitMock.EXPECT().CheckTraceBenefit(gomock.Any(), gomock.Any()).Return(nil, assert.AnError) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("").AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) - rateLimiterFactoryMock := limitermocks.NewMockIRateLimiter(ctrl) - rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() + rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) - rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), "query_trace:qps:space_id:123", 1, gomock.Any()).Return(nil, assert.AnError) return fields{ traceService: traceServiceMock, auth: authMock, @@ -2526,26 +3259,35 @@ func TestOpenAPIApplication_AllowBySpace(t *testing.T) { } }, args: args{ - ctx: context.Background(), - workspaceID: 123, + ctx: context.Background(), + req: &openapi.CreateAnnotationRequest{ + WorkspaceID: 1, + AnnotationValueType: ptr.Of(annotation.ValueType(loop_span.AnnotationValueTypeString)), + AnnotationValue: "test", + Base: &base.Base{Caller: "test"}, + }, }, - want: true, + want: nil, + wantErr: true, }, { - name: "nil result returns true", + name: "trace service fails", fieldsGetter: func(ctrl *gomock.Controller) fields { traceServiceMock := servicemocks.NewMockITraceService(ctrl) - authMock := rpcmocks.NewMockIAuthProvider(ctrl) + traceServiceMock.EXPECT().CreateAnnotation(gomock.Any(), gomock.Any()).Return(assert.AnError) benefitMock := benefitmocks.NewMockIBenefitService(ctrl) + benefitMock.EXPECT().CheckTraceBenefit(gomock.Any(), gomock.Any()).Return(&benefit.CheckTraceBenefitResult{ + StorageDuration: 3, + }, nil) + authMock := rpcmocks.NewMockIAuthProvider(ctrl) tenantMock := tenantmocks.NewMockITenantProvider(ctrl) workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("").AnyTimes() rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) - rateLimiterFactoryMock := limitermocks.NewMockIRateLimiter(ctrl) - rateLimiterMock.EXPECT().NewRateLimiter().Return(rateLimiterFactoryMock).AnyTimes() + rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() traceConfigMock := configmocks.NewMockITraceConfig(ctrl) metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) - traceConfigMock.EXPECT().GetQueryMaxQPSBySpace(gomock.Any(), int64(123)).Return(100, nil) - rateLimiterFactoryMock.EXPECT().AllowN(gomock.Any(), "query_trace:qps:space_id:123", 1, gomock.Any()).Return(nil, nil) return fields{ traceService: traceServiceMock, auth: authMock, @@ -2558,22 +3300,241 @@ func TestOpenAPIApplication_AllowBySpace(t *testing.T) { } }, args: args{ - ctx: context.Background(), - workspaceID: 123, + ctx: context.Background(), + req: &openapi.CreateAnnotationRequest{ + WorkspaceID: 1, + AnnotationValueType: ptr.Of(annotation.ValueType(loop_span.AnnotationValueTypeString)), + AnnotationValue: "test", + Base: &base.Base{Caller: "test"}, + }, }, - want: true, + want: nil, + wantErr: true, }, } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + fields := tt.fieldsGetter(ctrl) + o := &OpenAPIApplication{ + traceService: fields.traceService, + auth: fields.auth, + benefit: fields.benefit, + tenant: fields.tenant, + workspace: fields.workspace, + rateLimiter: fields.rateLimiter.NewRateLimiter(), + traceConfig: fields.traceConfig, + metrics: fields.metrics, + } + got, err := o.CreateAnnotation(tt.args.ctx, tt.args.req) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} +// 补充DeleteAnnotation的更多测试场景 +func TestOpenAPIApplication_DeleteAnnotation_AdditionalScenarios(t *testing.T) { + type fields struct { + traceService service.ITraceService + auth rpc.IAuthProvider + benefit benefit.IBenefitService + tenant tenant.ITenantProvider + workspace workspace.IWorkSpaceProvider + rateLimiter limiter.IRateLimiterFactory + traceConfig config.ITraceConfig + metrics metrics.ITraceMetrics + } + type args struct { + ctx context.Context + req *openapi.DeleteAnnotationRequest + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want *openapi.DeleteAnnotationResponse + wantErr bool + }{ + { + name: "benefit check fails", + fieldsGetter: func(ctrl *gomock.Controller) fields { + traceServiceMock := servicemocks.NewMockITraceService(ctrl) + authMock := rpcmocks.NewMockIAuthProvider(ctrl) + benefitMock := benefitmocks.NewMockIBenefitService(ctrl) + benefitMock.EXPECT().CheckTraceBenefit(gomock.Any(), gomock.Any()).Return(nil, assert.AnError) + tenantMock := tenantmocks.NewMockITenantProvider(ctrl) + workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("").AnyTimes() + rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) + rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() + traceConfigMock := configmocks.NewMockITraceConfig(ctrl) + metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) + return fields{ + traceService: traceServiceMock, + auth: authMock, + benefit: benefitMock, + tenant: tenantMock, + workspace: workspaceMock, + rateLimiter: rateLimiterMock, + traceConfig: traceConfigMock, + metrics: metricsMock, + } + }, + args: args{ + ctx: context.Background(), + req: &openapi.DeleteAnnotationRequest{ + WorkspaceID: 1, + Base: &base.Base{Caller: "test"}, + }, + }, + want: nil, + wantErr: true, + }, + { + name: "trace service fails", + fieldsGetter: func(ctrl *gomock.Controller) fields { + traceServiceMock := servicemocks.NewMockITraceService(ctrl) + traceServiceMock.EXPECT().DeleteAnnotation(gomock.Any(), gomock.Any()).Return(assert.AnError) + benefitMock := benefitmocks.NewMockIBenefitService(ctrl) + benefitMock.EXPECT().CheckTraceBenefit(gomock.Any(), gomock.Any()).Return(&benefit.CheckTraceBenefitResult{ + StorageDuration: 3, + }, nil) + authMock := rpcmocks.NewMockIAuthProvider(ctrl) + tenantMock := tenantmocks.NewMockITenantProvider(ctrl) + workspaceMock := workspacemocks.NewMockIWorkSpaceProvider(ctrl) + workspaceMock.EXPECT().GetThirdPartyQueryWorkSpaceID(gomock.Any(), int64(123)).Return("123").AnyTimes() + workspaceMock.EXPECT().GetIngestWorkSpaceID(gomock.Any(), gomock.Any()).Return("").AnyTimes() + rateLimiterMock := limitermocks.NewMockIRateLimiterFactory(ctrl) + rateLimiterMock.EXPECT().NewRateLimiter().Return(limitermocks.NewMockIRateLimiter(ctrl)).AnyTimes() + traceConfigMock := configmocks.NewMockITraceConfig(ctrl) + metricsMock := metricsmocks.NewMockITraceMetrics(ctrl) + return fields{ + traceService: traceServiceMock, + auth: authMock, + benefit: benefitMock, + tenant: tenantMock, + workspace: workspaceMock, + rateLimiter: rateLimiterMock, + traceConfig: traceConfigMock, + metrics: metricsMock, + } + }, + args: args{ + ctx: context.Background(), + req: &openapi.DeleteAnnotationRequest{ + WorkspaceID: 1, + Base: &base.Base{Caller: "test"}, + }, + }, + want: nil, + wantErr: true, + }, + } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { ctrl := gomock.NewController(t) defer ctrl.Finish() fields := tt.fieldsGetter(ctrl) - o, err := NewOpenAPIApplication(fields.traceService, fields.auth, fields.benefit, fields.tenant, fields.workspace, fields.rateLimiter, fields.traceConfig, fields.metrics) - assert.NoError(t, err) - got := o.(*OpenAPIApplication).AllowBySpace(tt.args.ctx, tt.args.workspaceID) + o := &OpenAPIApplication{ + traceService: fields.traceService, + auth: fields.auth, + benefit: fields.benefit, + tenant: fields.tenant, + workspace: fields.workspace, + rateLimiter: fields.rateLimiter.NewRateLimiter(), + traceConfig: fields.traceConfig, + metrics: fields.metrics, + } + got, err := o.DeleteAnnotation(tt.args.ctx, tt.args.req) + assert.Equal(t, tt.wantErr, err != nil) assert.Equal(t, tt.want, got) }) } } + +// 测试validate和build函数 +func TestOpenAPIApplication_validateIngestTracesReq(t *testing.T) { + app := &OpenAPIApplication{} + + // 测试nil请求 + err := app.validateIngestTracesReq(context.Background(), nil) + assert.Error(t, err) + + // 测试空spans + err = app.validateIngestTracesReq(context.Background(), &openapi.IngestTracesRequest{ + Spans: []*span.InputSpan{}, + }) + assert.Error(t, err) + + // 测试不同workspace id的spans + err = app.validateIngestTracesReq(context.Background(), &openapi.IngestTracesRequest{ + Spans: []*span.InputSpan{ + {WorkspaceID: "1"}, + {WorkspaceID: "2"}, + }, + }) + assert.Error(t, err) + + // 测试正常情况 + err = app.validateIngestTracesReq(context.Background(), &openapi.IngestTracesRequest{ + Spans: []*span.InputSpan{ + {WorkspaceID: "1"}, + {WorkspaceID: "1"}, + }, + }) + assert.NoError(t, err) +} + +func TestOpenAPIApplication_validateIngestTracesReqByTenant(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + traceConfigMock := configmocks.NewMockITraceConfig(ctrl) + app := &OpenAPIApplication{ + traceConfig: traceConfigMock, + } + + // 测试nil请求 + traceConfigMock.EXPECT().GetTraceIngestTenantProducerCfg(gomock.Any()).Return(nil, nil) + err := app.validateIngestTracesReqByTenant(context.Background(), "tenant", nil) + assert.Error(t, err) + + // 测试超过最大span长度 + traceConfigMock.EXPECT().GetTraceIngestTenantProducerCfg(gomock.Any()).Return(map[string]*config.IngestConfig{ + "tenant": {MaxSpanLength: 1}, + }, nil) + err = app.validateIngestTracesReqByTenant(context.Background(), "tenant", &openapi.IngestTracesRequest{ + Spans: []*span.InputSpan{{}, {}}, + }) + assert.Error(t, err) + + // 测试正常情况 + traceConfigMock.EXPECT().GetTraceIngestTenantProducerCfg(gomock.Any()).Return(nil, nil) + err = app.validateIngestTracesReqByTenant(context.Background(), "tenant", &openapi.IngestTracesRequest{ + Spans: []*span.InputSpan{{}}, + }) + assert.NoError(t, err) +} + +func TestOpenAPIApplication_unpackTenant(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + tenantMock := tenantmocks.NewMockITenantProvider(ctrl) + app := &OpenAPIApplication{ + tenant: tenantMock, + } + + // 测试nil spans + result := app.unpackTenant(context.Background(), nil) + assert.Nil(t, result) + + // 测试正常情况 + tenantMock.EXPECT().GetIngestTenant(gomock.Any(), gomock.Any()).Return("tenant1") + result = app.unpackTenant(context.Background(), []*loop_span.Span{{SpanID: "test"}}) + assert.Len(t, result, 1) + assert.Len(t, result["tenant1"], 1) +} diff --git a/backend/modules/observability/domain/component/config/config.go b/backend/modules/observability/domain/component/config/config.go index ad2d6a2c6..ea00b0efb 100644 --- a/backend/modules/observability/domain/component/config/config.go +++ b/backend/modules/observability/domain/component/config/config.go @@ -115,7 +115,7 @@ type ITraceConfig interface { GetTraceDataMaxDurationDay(ctx context.Context, platformType *string) int64 GetDefaultTraceTenant(ctx context.Context) string GetAnnotationSourceCfg(ctx context.Context) (*AnnotationSourceConfig, error) - GetQueryMaxQPSBySpace(ctx context.Context, workspaceID int64) (int, error) + GetQueryMaxQPS(ctx context.Context, key string) (int, error) conf.IConfigLoader } diff --git a/backend/modules/observability/domain/component/config/mocks/config.go b/backend/modules/observability/domain/component/config/mocks/config.go index b63baaf06..3e2d45000 100644 --- a/backend/modules/observability/domain/component/config/mocks/config.go +++ b/backend/modules/observability/domain/component/config/mocks/config.go @@ -130,19 +130,19 @@ func (mr *MockITraceConfigMockRecorder) GetPlatformTenants(ctx any) *gomock.Call return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPlatformTenants", reflect.TypeOf((*MockITraceConfig)(nil).GetPlatformTenants), ctx) } -// GetQueryMaxQPSBySpace mocks base method. -func (m *MockITraceConfig) GetQueryMaxQPSBySpace(ctx context.Context, workspaceID int64) (int, error) { +// GetQueryMaxQPS mocks base method. +func (m *MockITraceConfig) GetQueryMaxQPS(ctx context.Context, key string) (int, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetQueryMaxQPSBySpace", ctx, workspaceID) + ret := m.ctrl.Call(m, "GetQueryMaxQPS", ctx, key) ret0, _ := ret[0].(int) ret1, _ := ret[1].(error) return ret0, ret1 } -// GetQueryMaxQPSBySpace indicates an expected call of GetQueryMaxQPSBySpace. -func (mr *MockITraceConfigMockRecorder) GetQueryMaxQPSBySpace(ctx, workspaceID any) *gomock.Call { +// GetQueryMaxQPS indicates an expected call of GetQueryMaxQPS. +func (mr *MockITraceConfigMockRecorder) GetQueryMaxQPS(ctx, key any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetQueryMaxQPSBySpace", reflect.TypeOf((*MockITraceConfig)(nil).GetQueryMaxQPSBySpace), ctx, workspaceID) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetQueryMaxQPS", reflect.TypeOf((*MockITraceConfig)(nil).GetQueryMaxQPS), ctx, key) } // GetSystemViews mocks base method. diff --git a/backend/modules/observability/domain/component/metrics/metrics.go b/backend/modules/observability/domain/component/metrics/metrics.go index 186f12242..f4cba9d42 100644 --- a/backend/modules/observability/domain/component/metrics/metrics.go +++ b/backend/modules/observability/domain/component/metrics/metrics.go @@ -9,7 +9,5 @@ import "time" type ITraceMetrics interface { EmitListSpans(workspaceId int64, spanType string, start time.Time, isError bool) EmitGetTrace(workspaceId int64, start time.Time, isError bool) - EmitListSpansOapi(workspaceId int64, platformType, spanType string, spanSize int64, errorCode int, start time.Time, isError bool) - EmitSearchTraceOapi(workspaceId int64, platformType string, spanSize int64, errorCode int, start time.Time, isError bool) - EmitListTracesOapi(workspaceId int64, errorCode int, start time.Time, isError bool) + EmitTraceOapi(method string, workspaceId int64, platformType, spanListType string, spanSize int64, errorCode int, start time.Time, isError bool) } diff --git a/backend/modules/observability/domain/component/metrics/mocks/metrics.go b/backend/modules/observability/domain/component/metrics/mocks/metrics.go index 1010b7ac3..3bf3770b8 100644 --- a/backend/modules/observability/domain/component/metrics/mocks/metrics.go +++ b/backend/modules/observability/domain/component/metrics/mocks/metrics.go @@ -64,38 +64,14 @@ func (mr *MockITraceMetricsMockRecorder) EmitListSpans(workspaceId, spanType, st return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EmitListSpans", reflect.TypeOf((*MockITraceMetrics)(nil).EmitListSpans), workspaceId, spanType, start, isError) } -// EmitListSpansOapi mocks base method. -func (m *MockITraceMetrics) EmitListSpansOapi(workspaceId int64, platformType, spanType string, spanSize int64, errorCode int, start time.Time, isError bool) { +// EmitTraceOapi mocks base method. +func (m *MockITraceMetrics) EmitTraceOapi(method string, workspaceId int64, platformType, spanListType string, spanSize int64, errorCode int, start time.Time, isError bool) { m.ctrl.T.Helper() - m.ctrl.Call(m, "EmitListSpansOapi", workspaceId, platformType, spanType, spanSize, errorCode, start, isError) + m.ctrl.Call(m, "EmitTraceOapi", method, workspaceId, platformType, spanListType, spanSize, errorCode, start, isError) } -// EmitListSpansOapi indicates an expected call of EmitListSpansOapi. -func (mr *MockITraceMetricsMockRecorder) EmitListSpansOapi(workspaceId, platformType, spanType, spanSize, errorCode, start, isError any) *gomock.Call { +// EmitTraceOapi indicates an expected call of EmitTraceOapi. +func (mr *MockITraceMetricsMockRecorder) EmitTraceOapi(method, workspaceId, platformType, spanType, spanSize, errorCode, start, isError any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EmitListSpansOapi", reflect.TypeOf((*MockITraceMetrics)(nil).EmitListSpansOapi), workspaceId, platformType, spanType, spanSize, errorCode, start, isError) -} - -// EmitListTracesOapi mocks base method. -func (m *MockITraceMetrics) EmitListTracesOapi(workspaceId int64, errorCode int, start time.Time, isError bool) { - m.ctrl.T.Helper() - m.ctrl.Call(m, "EmitListTracesOapi", workspaceId, errorCode, start, isError) -} - -// EmitListTracesOapi indicates an expected call of EmitListTracesOapi. -func (mr *MockITraceMetricsMockRecorder) EmitListTracesOapi(workspaceId, errorCode, start, isError any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EmitListTracesOapi", reflect.TypeOf((*MockITraceMetrics)(nil).EmitListTracesOapi), workspaceId, errorCode, start, isError) -} - -// EmitSearchTraceOapi mocks base method. -func (m *MockITraceMetrics) EmitSearchTraceOapi(workspaceId int64, platformType string, spanSize int64, errorCode int, start time.Time, isError bool) { - m.ctrl.T.Helper() - m.ctrl.Call(m, "EmitSearchTraceOapi", workspaceId, platformType, spanSize, errorCode, start, isError) -} - -// EmitSearchTraceOapi indicates an expected call of EmitSearchTraceOapi. -func (mr *MockITraceMetricsMockRecorder) EmitSearchTraceOapi(workspaceId, platformType, spanSize, errorCode, start, isError any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EmitSearchTraceOapi", reflect.TypeOf((*MockITraceMetrics)(nil).EmitSearchTraceOapi), workspaceId, platformType, spanSize, errorCode, start, isError) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "EmitTraceOapi", reflect.TypeOf((*MockITraceMetrics)(nil).EmitTraceOapi), method, workspaceId, platformType, spanType, spanSize, errorCode, start, isError) } diff --git a/backend/modules/observability/domain/component/rpc/auth.go b/backend/modules/observability/domain/component/rpc/auth.go index 6340f78f7..9f9ffcb7a 100644 --- a/backend/modules/observability/domain/component/rpc/auth.go +++ b/backend/modules/observability/domain/component/rpc/auth.go @@ -8,6 +8,7 @@ import "context" const ( AuthActionTraceRead = "readLoopTrace" AuthActionTraceIngest = "ingestLoopTrace" + AuthActionTraceList = "listLoopTrace" AuthActionTraceViewCreate = "createLoopTraceView" AuthActionTraceViewList = "listLoopTraceView" AuthActionTraceViewEdit = "edit" diff --git a/backend/modules/observability/domain/component/workspace/mocks/workspace_provider.go b/backend/modules/observability/domain/component/workspace/mocks/workspace_provider.go index e753c2a1b..3e6496aa5 100644 --- a/backend/modules/observability/domain/component/workspace/mocks/workspace_provider.go +++ b/backend/modules/observability/domain/component/workspace/mocks/workspace_provider.go @@ -55,16 +55,16 @@ func (mr *MockIWorkSpaceProviderMockRecorder) GetIngestWorkSpaceID(ctx, spans an return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetIngestWorkSpaceID", reflect.TypeOf((*MockIWorkSpaceProvider)(nil).GetIngestWorkSpaceID), ctx, spans) } -// GetQueryWorkSpaceID mocks base method. -func (m *MockIWorkSpaceProvider) GetQueryWorkSpaceID(ctx context.Context, requestWorkspaceID int64) int64 { +// GetThirdPartyQueryWorkSpaceID mocks base method. +func (m *MockIWorkSpaceProvider) GetThirdPartyQueryWorkSpaceID(ctx context.Context, requestWorkspaceID int64) string { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetQueryWorkSpaceID", ctx, requestWorkspaceID) - ret0, _ := ret[0].(int64) + ret := m.ctrl.Call(m, "GetThirdPartyQueryWorkSpaceID", ctx, requestWorkspaceID) + ret0, _ := ret[0].(string) return ret0 } -// GetQueryWorkSpaceID indicates an expected call of GetQueryWorkSpaceID. -func (mr *MockIWorkSpaceProviderMockRecorder) GetQueryWorkSpaceID(ctx, requestWorkspaceID any) *gomock.Call { +// GetThirdPartyQueryWorkSpaceID indicates an expected call of GetThirdPartyQueryWorkSpaceID. +func (mr *MockIWorkSpaceProviderMockRecorder) GetThirdPartyQueryWorkSpaceID(ctx, requestWorkspaceID any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetQueryWorkSpaceID", reflect.TypeOf((*MockIWorkSpaceProvider)(nil).GetQueryWorkSpaceID), ctx, requestWorkspaceID) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetThirdPartyQueryWorkSpaceID", reflect.TypeOf((*MockIWorkSpaceProvider)(nil).GetThirdPartyQueryWorkSpaceID), ctx, requestWorkspaceID) } diff --git a/backend/modules/observability/domain/component/workspace/workspace.go b/backend/modules/observability/domain/component/workspace/workspace.go index 50b7bc4dc..b7b539124 100644 --- a/backend/modules/observability/domain/component/workspace/workspace.go +++ b/backend/modules/observability/domain/component/workspace/workspace.go @@ -12,5 +12,5 @@ import ( //go:generate mockgen -destination=mocks/workspace_provider.go -package=mocks . IWorkSpaceProvider type IWorkSpaceProvider interface { GetIngestWorkSpaceID(ctx context.Context, spans []*span.InputSpan) string - GetQueryWorkSpaceID(ctx context.Context, requestWorkspaceID int64) int64 + GetThirdPartyQueryWorkSpaceID(ctx context.Context, requestWorkspaceID int64) string } diff --git a/backend/modules/observability/domain/trace/service/trace/span_filter/cozeloop_filter.go b/backend/modules/observability/domain/trace/service/trace/span_filter/cozeloop_filter.go index e50d4e290..e65229754 100644 --- a/backend/modules/observability/domain/trace/service/trace/span_filter/cozeloop_filter.go +++ b/backend/modules/observability/domain/trace/service/trace/span_filter/cozeloop_filter.go @@ -18,7 +18,7 @@ func (c *CozeLoopFilter) BuildBasicSpanFilter(ctx context.Context, env *SpanEnv) { FieldName: loop_span.SpanFieldSpaceId, FieldType: loop_span.FieldTypeString, - Values: []string{strconv.FormatInt(env.WorkspaceId, 10)}, + Values: []string{strconv.FormatInt(env.WorkspaceID, 10)}, QueryType: ptr.Of(loop_span.QueryTypeEnumIn), }, { diff --git a/backend/modules/observability/domain/trace/service/trace/span_filter/cozeloop_filter_test.go b/backend/modules/observability/domain/trace/service/trace/span_filter/cozeloop_filter_test.go index fee04c1a9..2a379830a 100644 --- a/backend/modules/observability/domain/trace/service/trace/span_filter/cozeloop_filter_test.go +++ b/backend/modules/observability/domain/trace/service/trace/span_filter/cozeloop_filter_test.go @@ -21,7 +21,7 @@ func TestCozeLoopFilter_BuildBasicSpanFilter(t *testing.T) { }{ { name: "success", - env: &SpanEnv{WorkspaceId: 123}, + env: &SpanEnv{WorkspaceID: 123}, want: []*loop_span.FilterField{ { FieldName: loop_span.SpanFieldSpaceId, @@ -56,7 +56,7 @@ func TestCozeLoopFilter_BuildAllSpanFilter(t *testing.T) { }{ { name: "success", - env: &SpanEnv{WorkspaceId: 123}, + env: &SpanEnv{WorkspaceID: 123}, want: nil, }, } @@ -78,7 +78,7 @@ func TestCozeLoopFilter_BuildRootSpanFilter(t *testing.T) { }{ { name: "success", - env: &SpanEnv{WorkspaceId: 123}, + env: &SpanEnv{WorkspaceID: 123}, want: []*loop_span.FilterField{ { FieldName: loop_span.SpanFieldParentID, @@ -107,7 +107,7 @@ func TestCozeLoopFilter_BuildLlmSpanFilter(t *testing.T) { }{ { name: "success", - env: &SpanEnv{WorkspaceId: 123}, + env: &SpanEnv{WorkspaceID: 123}, want: []*loop_span.FilterField{ { FieldName: loop_span.SpanFieldSpanType, diff --git a/backend/modules/observability/domain/trace/service/trace/span_filter/eval_target_filter.go b/backend/modules/observability/domain/trace/service/trace/span_filter/eval_target_filter.go index 6e1af4429..a17ef1ddf 100644 --- a/backend/modules/observability/domain/trace/service/trace/span_filter/eval_target_filter.go +++ b/backend/modules/observability/domain/trace/service/trace/span_filter/eval_target_filter.go @@ -18,7 +18,7 @@ func (e *EvalTargetFilter) BuildBasicSpanFilter(ctx context.Context, env *SpanEn { FieldName: loop_span.SpanFieldSpaceId, FieldType: loop_span.FieldTypeString, - Values: []string{strconv.FormatInt(env.WorkspaceId, 10)}, + Values: []string{strconv.FormatInt(env.WorkspaceID, 10)}, QueryType: ptr.Of(loop_span.QueryTypeEnumIn), }, { diff --git a/backend/modules/observability/domain/trace/service/trace/span_filter/eval_target_filter_test.go b/backend/modules/observability/domain/trace/service/trace/span_filter/eval_target_filter_test.go index afbdc9a03..a3bd5942a 100644 --- a/backend/modules/observability/domain/trace/service/trace/span_filter/eval_target_filter_test.go +++ b/backend/modules/observability/domain/trace/service/trace/span_filter/eval_target_filter_test.go @@ -21,7 +21,7 @@ func TestEvalTargetFilter_BuildBasicSpanFilter(t *testing.T) { }{ { name: "success", - env: &SpanEnv{WorkspaceId: 123}, + env: &SpanEnv{WorkspaceID: 123}, want: []*loop_span.FilterField{ { FieldName: loop_span.SpanFieldSpaceId, @@ -56,7 +56,7 @@ func TestEvalTargetFilter_BuildRootSpanFilter(t *testing.T) { }{ { name: "success", - env: &SpanEnv{WorkspaceId: 123}, + env: &SpanEnv{WorkspaceID: 123}, want: []*loop_span.FilterField{ { FieldName: loop_span.SpanFieldParentID, @@ -85,7 +85,7 @@ func TestEvalTargetFilter_BuildLLMSpanFilter(t *testing.T) { }{ { name: "success", - env: &SpanEnv{WorkspaceId: 123}, + env: &SpanEnv{WorkspaceID: 123}, want: []*loop_span.FilterField{ { FieldName: loop_span.SpanFieldSpanType, @@ -114,7 +114,7 @@ func TestEvalTargetFilter_BuildALLSpanFilter(t *testing.T) { }{ { name: "success", - env: &SpanEnv{WorkspaceId: 123}, + env: &SpanEnv{WorkspaceID: 123}, want: nil, }, } diff --git a/backend/modules/observability/domain/trace/service/trace/span_filter/evaluator_filter.go b/backend/modules/observability/domain/trace/service/trace/span_filter/evaluator_filter.go index bff353e3b..32a4b12d4 100644 --- a/backend/modules/observability/domain/trace/service/trace/span_filter/evaluator_filter.go +++ b/backend/modules/observability/domain/trace/service/trace/span_filter/evaluator_filter.go @@ -18,7 +18,7 @@ func (e *EvaluatorFilter) BuildBasicSpanFilter(ctx context.Context, env *SpanEnv { FieldName: loop_span.SpanFieldSpaceId, FieldType: loop_span.FieldTypeString, - Values: []string{strconv.FormatInt(env.WorkspaceId, 10)}, + Values: []string{strconv.FormatInt(env.WorkspaceID, 10)}, QueryType: ptr.Of(loop_span.QueryTypeEnumIn), }, { diff --git a/backend/modules/observability/domain/trace/service/trace/span_filter/evaluator_filter_test.go b/backend/modules/observability/domain/trace/service/trace/span_filter/evaluator_filter_test.go index 16d9908fa..463ee1873 100644 --- a/backend/modules/observability/domain/trace/service/trace/span_filter/evaluator_filter_test.go +++ b/backend/modules/observability/domain/trace/service/trace/span_filter/evaluator_filter_test.go @@ -21,7 +21,7 @@ func TestEvaluatorFilter_BuildBasicSpanFilter(t *testing.T) { }{ { name: "success", - env: &SpanEnv{WorkspaceId: 123}, + env: &SpanEnv{WorkspaceID: 123}, want: []*loop_span.FilterField{ { FieldName: loop_span.SpanFieldSpaceId, diff --git a/backend/modules/observability/domain/trace/service/trace/span_filter/filter.go b/backend/modules/observability/domain/trace/service/trace/span_filter/filter.go index 34a5d542b..5e259a998 100644 --- a/backend/modules/observability/domain/trace/service/trace/span_filter/filter.go +++ b/backend/modules/observability/domain/trace/service/trace/span_filter/filter.go @@ -12,7 +12,8 @@ import ( ) type SpanEnv struct { - WorkspaceId int64 + WorkspaceID int64 + ThirdPartyWorkspaceID string } type Factory interface { diff --git a/backend/modules/observability/domain/trace/service/trace/span_filter/prompt_filter.go b/backend/modules/observability/domain/trace/service/trace/span_filter/prompt_filter.go index 0aa0350be..37d1bbce7 100644 --- a/backend/modules/observability/domain/trace/service/trace/span_filter/prompt_filter.go +++ b/backend/modules/observability/domain/trace/service/trace/span_filter/prompt_filter.go @@ -22,7 +22,7 @@ func (p *PromptFilter) BuildBasicSpanFilter(ctx context.Context, env *SpanEnv) ( { FieldName: loop_span.SpanFieldSpaceId, FieldType: loop_span.FieldTypeString, - Values: []string{strconv.FormatInt(env.WorkspaceId, 10)}, + Values: []string{strconv.FormatInt(env.WorkspaceID, 10)}, QueryType: ptr.Of(loop_span.QueryTypeEnumIn), }, { diff --git a/backend/modules/observability/domain/trace/service/trace/span_filter/prompt_filter_test.go b/backend/modules/observability/domain/trace/service/trace/span_filter/prompt_filter_test.go index 8eb8ee4a9..452e6211f 100644 --- a/backend/modules/observability/domain/trace/service/trace/span_filter/prompt_filter_test.go +++ b/backend/modules/observability/domain/trace/service/trace/span_filter/prompt_filter_test.go @@ -24,7 +24,7 @@ func TestPromptFilter_BuildBasicSpanFilter(t *testing.T) { }{ { name: "success", - env: &SpanEnv{WorkspaceId: 123}, + env: &SpanEnv{WorkspaceID: 123}, want: []*loop_span.FilterField{ { FieldName: loop_span.SpanFieldSpaceId, diff --git a/backend/modules/observability/domain/trace/service/trace/span_processor/processor.go b/backend/modules/observability/domain/trace/service/trace/span_processor/processor.go index cd9ea0912..465542a9c 100644 --- a/backend/modules/observability/domain/trace/service/trace/span_processor/processor.go +++ b/backend/modules/observability/domain/trace/service/trace/span_processor/processor.go @@ -11,11 +11,12 @@ import ( type Settings struct { // query parameters - WorkspaceId int64 - PlatformType loop_span.PlatformType - QueryStartTime int64 // ms - QueryEndTime int64 // ms - Tenant string + WorkspaceId int64 + ThirdPartyWorkspaceID string + PlatformType loop_span.PlatformType + QueryStartTime int64 // ms + QueryEndTime int64 // ms + Tenant string } type Factory interface { diff --git a/backend/modules/observability/domain/trace/service/trace_service.go b/backend/modules/observability/domain/trace/service/trace_service.go index b37f12ce8..b5e235ca7 100644 --- a/backend/modules/observability/domain/trace/service/trace_service.go +++ b/backend/modules/observability/domain/trace/service/trace_service.go @@ -31,15 +31,16 @@ import ( ) type ListSpansReq struct { - WorkspaceID int64 - StartTime int64 // ms - EndTime int64 // ms - Filters *loop_span.FilterFields - Limit int32 - DescByStartTime bool - PageToken string - PlatformType loop_span.PlatformType - SpanListType loop_span.SpanListType + WorkspaceID int64 + ThirdPartyWorkspaceID string + StartTime int64 // ms + EndTime int64 // ms + Filters *loop_span.FilterFields + Limit int32 + DescByStartTime bool + PageToken string + PlatformType loop_span.PlatformType + SpanListType loop_span.SpanListType } type ListSpansResp struct { @@ -64,14 +65,15 @@ type GetTraceResp struct { } type SearchTraceOApiReq struct { - WorkspaceID int64 - Tenants []string - TraceID string - LogID string - StartTime int64 // ms - EndTime int64 // ms - Limit int32 - PlatformType loop_span.PlatformType + WorkspaceID int64 + ThirdPartyWorkspaceID string + Tenants []string + TraceID string + LogID string + StartTime int64 // ms + EndTime int64 // ms + Limit int32 + PlatformType loop_span.PlatformType } type SearchTraceOApiResp struct { @@ -79,16 +81,17 @@ type SearchTraceOApiResp struct { } type ListSpansOApiReq struct { - WorkspaceID int64 - Tenants []string - StartTime int64 // ms - EndTime int64 // ms - Filters *loop_span.FilterFields - Limit int32 - DescByStartTime bool - PageToken string - PlatformType loop_span.PlatformType - SpanListType loop_span.SpanListType + WorkspaceID int64 + ThirdPartyWorkspaceID string + Tenants []string + StartTime int64 // ms + EndTime int64 // ms + Filters *loop_span.FilterFields + Limit int32 + DescByStartTime bool + PageToken string + PlatformType loop_span.PlatformType + SpanListType loop_span.SpanListType } type ListSpansOApiResp struct { @@ -104,9 +107,10 @@ type TraceQueryParam struct { } type GetTracesAdvanceInfoReq struct { - WorkspaceID int64 - Traces []*TraceQueryParam - PlatformType loop_span.PlatformType + WorkspaceID int64 + ThirdPartyWorkspaceID string + Traces []*TraceQueryParam + PlatformType loop_span.PlatformType } type GetTracesAdvanceInfoResp struct { @@ -354,10 +358,11 @@ func (r *TraceServiceImpl) SearchTraceOApi(ctx context.Context, req *SearchTrace return nil, err } processors, err := r.buildHelper.BuildSearchTraceOApiProcessors(ctx, span_processor.Settings{ - WorkspaceId: req.WorkspaceID, - QueryStartTime: req.StartTime, - QueryEndTime: req.EndTime, - PlatformType: req.PlatformType, + WorkspaceId: req.WorkspaceID, + ThirdPartyWorkspaceID: req.ThirdPartyWorkspaceID, + QueryStartTime: req.StartTime, + QueryEndTime: req.EndTime, + PlatformType: req.PlatformType, }) if err != nil { return nil, errorx.WrapByCode(err, obErrorx.CommercialCommonInternalErrorCodeCode) @@ -383,8 +388,9 @@ func (r *TraceServiceImpl) ListSpansOApi(ctx context.Context, req *ListSpansOApi return nil, err } builtinFilter, err := r.buildBuiltinFilters(ctx, platformFilter, &ListSpansReq{ - WorkspaceID: req.WorkspaceID, - SpanListType: req.SpanListType, + WorkspaceID: req.WorkspaceID, + ThirdPartyWorkspaceID: req.ThirdPartyWorkspaceID, + SpanListType: req.SpanListType, }) if err != nil { return nil, err @@ -920,7 +926,8 @@ func (r *TraceServiceImpl) getAnnotationCallerCfg(ctx context.Context, caller st func (r *TraceServiceImpl) buildBuiltinFilters(ctx context.Context, f span_filter.Filter, req *ListSpansReq) (*loop_span.FilterFields, error) { filters := make([]*loop_span.FilterField, 0) env := &span_filter.SpanEnv{ - WorkspaceId: req.WorkspaceID, + WorkspaceID: req.WorkspaceID, + ThirdPartyWorkspaceID: req.ThirdPartyWorkspaceID, } basicFilter, forceQuery, err := f.BuildBasicSpanFilter(ctx, env) if err != nil { diff --git a/backend/modules/observability/domain/trace/service/trace_service_test.go b/backend/modules/observability/domain/trace/service/trace_service_test.go index 23f30a0d9..c38827f30 100644 --- a/backend/modules/observability/domain/trace/service/trace_service_test.go +++ b/backend/modules/observability/domain/trace/service/trace_service_test.go @@ -2532,6 +2532,303 @@ func TestTraceServiceImpl_ListSpansOApi(t *testing.T) { want *ListSpansOApiResp wantErr bool }{ + { + name: "list spans oapi successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + repoMock := repomocks.NewMockITraceRepo(ctrl) + repoMock.EXPECT().ListSpans(gomock.Any(), gomock.Any()).Return(&repo.ListSpansResult{ + Spans: loop_span.SpanList{ + { + TraceID: "trace-123", + SpanID: "span-456", + }, + }, + PageToken: "next-token", + HasMore: true, + }, nil) + + filterFactoryMock := filtermocks.NewMockPlatformFilterFactory(ctrl) + filterMock := filtermocks.NewMockFilter(ctrl) + filterMock.EXPECT().BuildBasicSpanFilter(gomock.Any(), gomock.Any()).Return([]*loop_span.FilterField{ + { + FieldName: loop_span.SpanFieldSpaceId, + FieldType: loop_span.FieldTypeString, + Values: []string{"123"}, + QueryType: ptr.Of(loop_span.QueryTypeEnumIn), + }, + }, false, nil) + filterMock.EXPECT().BuildALLSpanFilter(gomock.Any(), gomock.Any()).Return(nil, nil) + filterFactoryMock.EXPECT().GetFilter(gomock.Any(), gomock.Any()).Return(filterMock, nil) + + buildHelper := NewTraceFilterProcessorBuilder(filterFactoryMock, nil, nil, nil, nil, nil, nil) + + return fields{ + traceRepo: repoMock, + buildHelper: buildHelper, + } + }, + args: args{ + ctx: context.Background(), + req: &ListSpansOApiReq{ + WorkspaceID: 123, + Tenants: []string{"tenant1"}, + StartTime: 1640995200000, + EndTime: 1640995800000, + Limit: 100, + PlatformType: loop_span.PlatformCozeLoop, + SpanListType: loop_span.SpanListTypeAllSpan, + }, + }, + want: &ListSpansOApiResp{ + Spans: loop_span.SpanList{ + { + TraceID: "trace-123", + SpanID: "span-456", + }, + }, + NextPageToken: "next-token", + HasMore: true, + }, + wantErr: false, + }, + { + name: "list spans oapi with empty builtin filter", + fieldsGetter: func(ctrl *gomock.Controller) fields { + filterFactoryMock := filtermocks.NewMockPlatformFilterFactory(ctrl) + filterMock := filtermocks.NewMockFilter(ctrl) + filterMock.EXPECT().BuildBasicSpanFilter(gomock.Any(), gomock.Any()).Return([]*loop_span.FilterField{}, false, nil) + filterFactoryMock.EXPECT().GetFilter(gomock.Any(), gomock.Any()).Return(filterMock, nil) + + buildHelper := NewTraceFilterProcessorBuilder(filterFactoryMock, nil, nil, nil, nil, nil, nil) + + return fields{ + buildHelper: buildHelper, + } + }, + args: args{ + ctx: context.Background(), + req: &ListSpansOApiReq{ + WorkspaceID: 123, + Tenants: []string{"tenant1"}, + StartTime: 1640995200000, + EndTime: 1640995800000, + Limit: 100, + PlatformType: loop_span.PlatformCozeLoop, + SpanListType: loop_span.SpanListTypeAllSpan, + }, + }, + want: &ListSpansOApiResp{ + Spans: loop_span.SpanList{}, + }, + wantErr: false, + }, + { + name: "list spans oapi with multiple processors", + fieldsGetter: func(ctrl *gomock.Controller) fields { + repoMock := repomocks.NewMockITraceRepo(ctrl) + repoMock.EXPECT().ListSpans(gomock.Any(), gomock.Any()).Return(&repo.ListSpansResult{ + Spans: loop_span.SpanList{ + { + TraceID: "trace-123", + SpanID: "span-456", + WorkspaceID: "123", + }, + }, + PageToken: "", + HasMore: false, + }, nil) + + filterFactoryMock := filtermocks.NewMockPlatformFilterFactory(ctrl) + filterMock := filtermocks.NewMockFilter(ctrl) + filterMock.EXPECT().BuildBasicSpanFilter(gomock.Any(), gomock.Any()).Return([]*loop_span.FilterField{ + { + FieldName: loop_span.SpanFieldSpaceId, + FieldType: loop_span.FieldTypeString, + Values: []string{"123"}, + QueryType: ptr.Of(loop_span.QueryTypeEnumIn), + }, + }, false, nil) + filterMock.EXPECT().BuildALLSpanFilter(gomock.Any(), gomock.Any()).Return(nil, nil) + filterFactoryMock.EXPECT().GetFilter(gomock.Any(), gomock.Any()).Return(filterMock, nil) + + buildHelper := NewTraceFilterProcessorBuilder(filterFactoryMock, nil, nil, nil, nil, nil, + []span_processor.Factory{span_processor.NewCheckProcessorFactory()}) + + return fields{ + traceRepo: repoMock, + buildHelper: buildHelper, + } + }, + args: args{ + ctx: context.Background(), + req: &ListSpansOApiReq{ + WorkspaceID: 123, + Tenants: []string{"tenant1"}, + StartTime: 1640995200000, + EndTime: 1640995800000, + Limit: 100, + PlatformType: loop_span.PlatformCozeLoop, + SpanListType: loop_span.SpanListTypeAllSpan, + }, + }, + want: &ListSpansOApiResp{ + Spans: loop_span.SpanList{ + { + TraceID: "trace-123", + SpanID: "span-456", + WorkspaceID: "123", + }, + }, + NextPageToken: "", + HasMore: false, + }, + wantErr: false, + }, + { + name: "list spans oapi failed due to platform filter error", + fieldsGetter: func(ctrl *gomock.Controller) fields { + filterFactoryMock := filtermocks.NewMockPlatformFilterFactory(ctrl) + filterFactoryMock.EXPECT().GetFilter(gomock.Any(), gomock.Any()).Return(nil, fmt.Errorf("platform filter error")) + + buildHelper := NewTraceFilterProcessorBuilder(filterFactoryMock, nil, nil, nil, nil, nil, nil) + + return fields{ + buildHelper: buildHelper, + } + }, + args: args{ + ctx: context.Background(), + req: &ListSpansOApiReq{ + WorkspaceID: 123, + Tenants: []string{"tenant1"}, + StartTime: 1640995200000, + EndTime: 1640995800000, + Limit: 100, + PlatformType: loop_span.PlatformCozeLoop, + SpanListType: loop_span.SpanListTypeAllSpan, + }, + }, + wantErr: true, + }, + { + name: "list spans oapi failed due to builtin filter error", + fieldsGetter: func(ctrl *gomock.Controller) fields { + filterFactoryMock := filtermocks.NewMockPlatformFilterFactory(ctrl) + filterMock := filtermocks.NewMockFilter(ctrl) + filterMock.EXPECT().BuildBasicSpanFilter(gomock.Any(), gomock.Any()).Return(nil, false, fmt.Errorf("builtin filter error")) + filterFactoryMock.EXPECT().GetFilter(gomock.Any(), gomock.Any()).Return(filterMock, nil) + + buildHelper := NewTraceFilterProcessorBuilder(filterFactoryMock, nil, nil, nil, nil, nil, nil) + + return fields{ + buildHelper: buildHelper, + } + }, + args: args{ + ctx: context.Background(), + req: &ListSpansOApiReq{ + WorkspaceID: 123, + Tenants: []string{"tenant1"}, + StartTime: 1640995200000, + EndTime: 1640995800000, + Limit: 100, + PlatformType: loop_span.PlatformCozeLoop, + SpanListType: loop_span.SpanListTypeAllSpan, + }, + }, + wantErr: true, + }, + { + name: "list spans oapi failed due to repo error", + fieldsGetter: func(ctrl *gomock.Controller) fields { + repoMock := repomocks.NewMockITraceRepo(ctrl) + repoMock.EXPECT().ListSpans(gomock.Any(), gomock.Any()).Return(nil, fmt.Errorf("repo error")) + + filterFactoryMock := filtermocks.NewMockPlatformFilterFactory(ctrl) + filterMock := filtermocks.NewMockFilter(ctrl) + filterMock.EXPECT().BuildBasicSpanFilter(gomock.Any(), gomock.Any()).Return([]*loop_span.FilterField{ + { + FieldName: loop_span.SpanFieldSpaceId, + FieldType: loop_span.FieldTypeString, + Values: []string{"123"}, + QueryType: ptr.Of(loop_span.QueryTypeEnumIn), + }, + }, false, nil) + filterMock.EXPECT().BuildALLSpanFilter(gomock.Any(), gomock.Any()).Return(nil, nil) + filterFactoryMock.EXPECT().GetFilter(gomock.Any(), gomock.Any()).Return(filterMock, nil) + + buildHelper := NewTraceFilterProcessorBuilder(filterFactoryMock, nil, nil, nil, nil, nil, nil) + + return fields{ + traceRepo: repoMock, + buildHelper: buildHelper, + } + }, + args: args{ + ctx: context.Background(), + req: &ListSpansOApiReq{ + WorkspaceID: 123, + Tenants: []string{"tenant1"}, + StartTime: 1640995200000, + EndTime: 1640995800000, + Limit: 100, + PlatformType: loop_span.PlatformCozeLoop, + SpanListType: loop_span.SpanListTypeAllSpan, + }, + }, + wantErr: true, + }, + { + name: "list spans oapi failed due to processor transform error", + fieldsGetter: func(ctrl *gomock.Controller) fields { + repoMock := repomocks.NewMockITraceRepo(ctrl) + repoMock.EXPECT().ListSpans(gomock.Any(), gomock.Any()).Return(&repo.ListSpansResult{ + Spans: loop_span.SpanList{ + { + TraceID: "trace-123", + SpanID: "span-456", + WorkspaceID: "1234", + }, + }, + PageToken: "", + HasMore: false, + }, nil) + + filterFactoryMock := filtermocks.NewMockPlatformFilterFactory(ctrl) + filterMock := filtermocks.NewMockFilter(ctrl) + filterMock.EXPECT().BuildBasicSpanFilter(gomock.Any(), gomock.Any()).Return([]*loop_span.FilterField{ + { + FieldName: loop_span.SpanFieldSpaceId, + FieldType: loop_span.FieldTypeString, + Values: []string{"123"}, + QueryType: ptr.Of(loop_span.QueryTypeEnumIn), + }, + }, false, nil) + filterMock.EXPECT().BuildALLSpanFilter(gomock.Any(), gomock.Any()).Return(nil, nil) + filterFactoryMock.EXPECT().GetFilter(gomock.Any(), gomock.Any()).Return(filterMock, nil) + + buildHelper := NewTraceFilterProcessorBuilder(filterFactoryMock, nil, nil, nil, nil, nil, + []span_processor.Factory{span_processor.NewCheckProcessorFactory()}) + + return fields{ + traceRepo: repoMock, + buildHelper: buildHelper, + } + }, + args: args{ + ctx: context.Background(), + req: &ListSpansOApiReq{ + WorkspaceID: 123, + Tenants: []string{"tenant1"}, + StartTime: 1640995200000, + EndTime: 1640995800000, + Limit: 100, + PlatformType: loop_span.PlatformCozeLoop, + SpanListType: loop_span.SpanListTypeAllSpan, + }, + }, + wantErr: true, + }, { name: "list spans failed due to invalid filter", fieldsGetter: func(ctrl *gomock.Controller) fields { @@ -2584,3 +2881,159 @@ func TestTraceServiceImpl_ListSpansOApi(t *testing.T) { }) } } + +func TestTraceFilterProcessorBuilderImpl_BuildListSpansOApiProcessors(t *testing.T) { + tests := []struct { + name string + listSpansOApiProcessorFactories []span_processor.Factory + want int + wantErr bool + }{ + { + name: "build processors successfully with empty factories", + listSpansOApiProcessorFactories: []span_processor.Factory{}, + want: 0, + wantErr: false, + }, + { + name: "build processors successfully with multiple factories", + listSpansOApiProcessorFactories: []span_processor.Factory{ + span_processor.NewCheckProcessorFactory(), + span_processor.NewCheckProcessorFactory(), + }, + want: 2, + wantErr: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + filterFactoryMock := filtermocks.NewMockPlatformFilterFactory(ctrl) + builder := NewTraceFilterProcessorBuilder( + filterFactoryMock, + nil, + nil, + nil, + nil, + nil, + tt.listSpansOApiProcessorFactories, + ) + + got, err := builder.BuildListSpansOApiProcessors(context.Background(), span_processor.Settings{ + WorkspaceId: 123, + QueryStartTime: 1640995200000, + QueryEndTime: 1640995800000, + }) + + assert.Equal(t, tt.wantErr, err != nil) + if !tt.wantErr { + assert.Equal(t, tt.want, len(got)) + } + }) + } +} + +func TestTraceFilterProcessorBuilderImpl_BuildIngestTraceProcessors_ErrorHandling(t *testing.T) { + tests := []struct { + name string + ingestTraceProcessorFactories []span_processor.Factory + want int + wantErr bool + }{ + { + name: "build ingest processors successfully with empty factories", + ingestTraceProcessorFactories: []span_processor.Factory{}, + want: 0, + wantErr: false, + }, + { + name: "build ingest processors successfully with multiple factories", + ingestTraceProcessorFactories: []span_processor.Factory{ + span_processor.NewCheckProcessorFactory(), + }, + want: 1, + wantErr: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + filterFactoryMock := filtermocks.NewMockPlatformFilterFactory(ctrl) + + builder := NewTraceFilterProcessorBuilder( + filterFactoryMock, + nil, + nil, + nil, + tt.ingestTraceProcessorFactories, + nil, + nil, + ) + + got, err := builder.BuildIngestTraceProcessors(context.Background(), span_processor.Settings{ + Tenant: "test-tenant", + }) + + assert.Equal(t, tt.wantErr, err != nil) + if !tt.wantErr { + assert.Equal(t, tt.want, len(got)) + } + }) + } +} + +func TestTraceFilterProcessorBuilderImpl_BuildSearchTraceOApiProcessors_ErrorHandling(t *testing.T) { + tests := []struct { + name string + searchTraceOApiProcessorFactories []span_processor.Factory + want int + wantErr bool + }{ + { + name: "build search trace oapi processors successfully with empty factories", + searchTraceOApiProcessorFactories: []span_processor.Factory{}, + want: 0, + wantErr: false, + }, + { + name: "build search trace oapi processors successfully with multiple factories", + searchTraceOApiProcessorFactories: []span_processor.Factory{ + span_processor.NewCheckProcessorFactory(), + }, + want: 1, + wantErr: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + filterFactoryMock := filtermocks.NewMockPlatformFilterFactory(ctrl) + builder := NewTraceFilterProcessorBuilder( + filterFactoryMock, + nil, + nil, + nil, + nil, + tt.searchTraceOApiProcessorFactories, + nil, + ) + + got, err := builder.BuildSearchTraceOApiProcessors(context.Background(), span_processor.Settings{ + WorkspaceId: 123, + QueryStartTime: 1640995200000, + QueryEndTime: 1640995800000, + }) + + assert.Equal(t, tt.wantErr, err != nil) + if !tt.wantErr { + assert.Equal(t, tt.want, len(got)) + } + }) + } +} diff --git a/backend/modules/observability/infra/config/trace.go b/backend/modules/observability/infra/config/trace.go index 707e4df3d..9d8cd8ed0 100644 --- a/backend/modules/observability/infra/config/trace.go +++ b/backend/modules/observability/infra/config/trace.go @@ -6,7 +6,6 @@ package config import ( "context" "fmt" - "strconv" "github.com/coze-dev/coze-loop/backend/modules/observability/domain/component/config" "github.com/coze-dev/coze-loop/backend/pkg/conf" @@ -141,12 +140,12 @@ func (t *TraceConfigCenter) GetAnnotationSourceCfg(ctx context.Context) (*config return annotationSourceCfg, nil } -func (t *TraceConfigCenter) GetQueryMaxQPSBySpace(ctx context.Context, workspaceID int64) (int, error) { +func (t *TraceConfigCenter) GetQueryMaxQPS(ctx context.Context, key string) (int, error) { qpsConfig := new(config.QueryTraceRateLimitConfig) if err := t.UnmarshalKey(ctx, queryTraceRateLimitCfgKey, &qpsConfig); err != nil { return 0, err } - if qps, ok := qpsConfig.SpaceMaxQPS[strconv.FormatInt(workspaceID, 10)]; ok { + if qps, ok := qpsConfig.SpaceMaxQPS[key]; ok { return qps, nil } return qpsConfig.DefaultMaxQPS, nil diff --git a/backend/modules/observability/infra/config/trace_test.go b/backend/modules/observability/infra/config/trace_test.go new file mode 100755 index 000000000..999e53ba6 --- /dev/null +++ b/backend/modules/observability/infra/config/trace_test.go @@ -0,0 +1,963 @@ +// Copyright (c) 2025 coze-dev Authors +// SPDX-License-Identifier: Apache-2.0 + +package config + +import ( + "context" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/coze-dev/coze-loop/backend/modules/observability/domain/component/config" + "github.com/coze-dev/coze-loop/backend/modules/observability/domain/trace/entity/loop_span" + confmocks "github.com/coze-dev/coze-loop/backend/pkg/conf/mocks" +) + +func TestTraceConfigCenter_GetSystemViews(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + } + type args struct { + ctx context.Context + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want []*config.SystemView + wantErr bool + }{ + { + name: "get system views successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), systemViewsCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + views := v.(*[]*config.SystemView) + *views = []*config.SystemView{ + {ID: 1, ViewName: "View 1"}, + {ID: 2, ViewName: "View 2"}, + } + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: []*config.SystemView{ + {ID: 1, ViewName: "View 1"}, + {ID: 2, ViewName: "View 2"}, + }, + wantErr: false, + }, + { + name: "unmarshal key failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), systemViewsCfgKey, gomock.Any()). + Return(fmt.Errorf("unmarshal error")) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: nil, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + } + got, err := tr.GetSystemViews(tt.args.ctx) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_GetPlatformTenants(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + } + type args struct { + ctx context.Context + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want *config.PlatformTenantsCfg + wantErr bool + }{ + { + name: "get platform tenants successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), platformTenantCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(*config.PlatformTenantsCfg) + cfg.Config = map[string][]string{"platform1": {"tenant1"}} + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: &config.PlatformTenantsCfg{ + Config: map[string][]string{"platform1": {"tenant1"}}, + }, + wantErr: false, + }, + { + name: "unmarshal key failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), platformTenantCfgKey, gomock.Any()). + Return(fmt.Errorf("unmarshal error")) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: nil, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + } + got, err := tr.GetPlatformTenants(tt.args.ctx) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_GetPlatformSpansTrans(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + } + type args struct { + ctx context.Context + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want *config.SpanTransHandlerConfig + wantErr bool + }{ + { + name: "get platform spans trans successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), platformSpanHandlerCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(*config.SpanTransHandlerConfig) + cfg.PlatformCfg = make(map[string]loop_span.SpanTransCfgList) + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: &config.SpanTransHandlerConfig{ + PlatformCfg: make(map[string]loop_span.SpanTransCfgList), + }, + wantErr: false, + }, + { + name: "unmarshal key failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), platformSpanHandlerCfgKey, gomock.Any()). + Return(fmt.Errorf("unmarshal error")) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: nil, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + } + got, err := tr.GetPlatformSpansTrans(tt.args.ctx) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_GetTraceIngestTenantProducerCfg(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + } + type args struct { + ctx context.Context + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want map[string]*config.IngestConfig + wantErr bool + }{ + { + name: "get trace ingest tenant producer cfg successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(context.Background(), traceIngestTenantCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(*map[string]*config.IngestConfig) + (*cfg)["tenant1"] = &config.IngestConfig{ + MaxSpanLength: 1000, + MqProducer: config.MqProducerCfg{Topic: "topic1"}, + } + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: map[string]*config.IngestConfig{ + "tenant1": { + MaxSpanLength: 1000, + MqProducer: config.MqProducerCfg{Topic: "topic1"}, + }, + }, + wantErr: false, + }, + { + name: "unmarshal key failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(context.Background(), traceIngestTenantCfgKey, gomock.Any()). + Return(fmt.Errorf("unmarshal error")) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: nil, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + } + got, err := tr.GetTraceIngestTenantProducerCfg(tt.args.ctx) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_GetAnnotationMqProducerCfg(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + } + type args struct { + ctx context.Context + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want *config.MqProducerCfg + wantErr bool + }{ + { + name: "get annotation mq producer cfg successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(context.Background(), annotationMqProducerCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(*config.MqProducerCfg) + cfg.Topic = "annotation_topic" + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: &config.MqProducerCfg{ + Topic: "annotation_topic", + }, + wantErr: false, + }, + { + name: "unmarshal key failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(context.Background(), annotationMqProducerCfgKey, gomock.Any()). + Return(fmt.Errorf("unmarshal error")) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: nil, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + } + got, err := tr.GetAnnotationMqProducerCfg(tt.args.ctx) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_GetTraceCkCfg(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + } + type args struct { + ctx context.Context + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want *config.TraceCKCfg + wantErr bool + }{ + { + name: "get trace ck cfg successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(context.Background(), traceCkCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(*config.TraceCKCfg) + cfg.DataBase = "trace_db" + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: &config.TraceCKCfg{ + DataBase: "trace_db", + }, + wantErr: false, + }, + { + name: "unmarshal key failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(context.Background(), traceCkCfgKey, gomock.Any()). + Return(fmt.Errorf("unmarshal error")) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: nil, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + } + got, err := tr.GetTraceCkCfg(tt.args.ctx) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_GetTenantConfig(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + } + type args struct { + ctx context.Context + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want *config.TenantCfg + wantErr bool + }{ + { + name: "get tenant config successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), tenantTablesCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(**config.TenantCfg) + *cfg = &config.TenantCfg{ + DefaultIngestTenant: "default_tenant", + } + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: &config.TenantCfg{ + DefaultIngestTenant: "default_tenant", + }, + wantErr: false, + }, + { + name: "unmarshal key failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), tenantTablesCfgKey, gomock.Any()). + Return(fmt.Errorf("unmarshal error")) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: nil, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + } + got, err := tr.GetTenantConfig(tt.args.ctx) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_GetTraceFieldMetaInfo(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + } + type args struct { + ctx context.Context + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want *config.TraceFieldMetaInfoCfg + wantErr bool + }{ + { + name: "get trace field meta info successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), traceFieldMetaInfoCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(**config.TraceFieldMetaInfoCfg) + *cfg = &config.TraceFieldMetaInfoCfg{ + AvailableFields: map[string]*config.FieldMeta{"field1": {}}, + } + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: &config.TraceFieldMetaInfoCfg{ + AvailableFields: map[string]*config.FieldMeta{"field1": {}}, + }, + wantErr: false, + }, + { + name: "unmarshal key failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), traceFieldMetaInfoCfgKey, gomock.Any()). + Return(fmt.Errorf("unmarshal error")) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: nil, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + } + got, err := tr.GetTraceFieldMetaInfo(tt.args.ctx) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_GetTraceDataMaxDurationDay(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + } + type args struct { + ctx context.Context + platformPtr *string + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want int64 + }{ + { + name: "platform ptr is nil, return default duration", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background(), platformPtr: nil}, + want: 7, + }, + { + name: "get duration successfully with platform type", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), traceMaxDurationDay, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + mp := v.(*map[string]int64) + (*mp)["coze"] = 30 + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background(), platformPtr: stringPtr("coze")}, + want: 30, + }, + { + name: "platform type not found, return default duration", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), traceMaxDurationDay, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + mp := v.(*map[string]int64) + (*mp)["other"] = 15 + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background(), platformPtr: stringPtr("coze")}, + want: 7, + }, + { + name: "platform type has zero value, return default duration", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), traceMaxDurationDay, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + mp := v.(*map[string]int64) + (*mp)["coze"] = 0 + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background(), platformPtr: stringPtr("coze")}, + want: 7, + }, + { + name: "unmarshal key failed, return default duration", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), traceMaxDurationDay, gomock.Any()). + Return(fmt.Errorf("unmarshal error")) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background(), platformPtr: stringPtr("coze")}, + want: 7, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + } + got := tr.GetTraceDataMaxDurationDay(tt.args.ctx, tt.args.platformPtr) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_GetDefaultTraceTenant(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + traceDefaultTenant string + } + type args struct { + ctx context.Context + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want string + }{ + { + name: "get default trace tenant successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + return fields{ + configLoader: mockLoader, + traceDefaultTenant: "default_tenant", + } + }, + args: args{ctx: context.Background()}, + want: "default_tenant", + }, + { + name: "get empty default trace tenant", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + return fields{ + configLoader: mockLoader, + traceDefaultTenant: "", + } + }, + args: args{ctx: context.Background()}, + want: "", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + traceDefaultTenant: f.traceDefaultTenant, + } + got := tr.GetDefaultTraceTenant(tt.args.ctx) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_getDefaultTraceTenant(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + traceDefaultTenant string + } + type args struct { + ctx context.Context + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want string + wantErr bool + }{ + { + name: "trace default tenant already set", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + return fields{ + configLoader: mockLoader, + traceDefaultTenant: "existing_tenant", + } + }, + args: args{ctx: context.Background()}, + want: "existing_tenant", + wantErr: false, + }, + { + name: "get tenant config successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), tenantTablesCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(**config.TenantCfg) + *cfg = &config.TenantCfg{ + DefaultIngestTenant: "new_tenant", + } + return nil + }) + return fields{ + configLoader: mockLoader, + traceDefaultTenant: "", + } + }, + args: args{ctx: context.Background()}, + want: "new_tenant", + wantErr: false, + }, + { + name: "get tenant config failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), tenantTablesCfgKey, gomock.Any()). + Return(fmt.Errorf("config error")) + return fields{ + configLoader: mockLoader, + traceDefaultTenant: "", + } + }, + args: args{ctx: context.Background()}, + want: "", + wantErr: true, + }, + { + name: "default ingest tenant is empty", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), tenantTablesCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(**config.TenantCfg) + *cfg = &config.TenantCfg{ + DefaultIngestTenant: "", + } + return nil + }) + return fields{ + configLoader: mockLoader, + traceDefaultTenant: "", + } + }, + args: args{ctx: context.Background()}, + want: "", + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + traceDefaultTenant: f.traceDefaultTenant, + } + got, err := tr.getDefaultTraceTenant(tt.args.ctx) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_GetAnnotationSourceCfg(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + } + type args struct { + ctx context.Context + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want *config.AnnotationSourceConfig + wantErr bool + }{ + { + name: "get annotation source cfg successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), annotationSourceCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(**config.AnnotationSourceConfig) + *cfg = &config.AnnotationSourceConfig{ + SourceCfg: map[string]config.AnnotationConfig{"source1": {}}, + } + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: &config.AnnotationSourceConfig{ + SourceCfg: map[string]config.AnnotationConfig{"source1": {}}, + }, + wantErr: false, + }, + { + name: "unmarshal key failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), annotationSourceCfgKey, gomock.Any()). + Return(fmt.Errorf("unmarshal error")) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background()}, + want: nil, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + } + got, err := tr.GetAnnotationSourceCfg(tt.args.ctx) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestTraceConfigCenter_GetQueryMaxQPS(t *testing.T) { + type fields struct { + configLoader *confmocks.MockIConfigLoader + } + type args struct { + ctx context.Context + key string + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + want int + wantErr bool + }{ + { + name: "get query max qps successfully with specific key", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), queryTraceRateLimitCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(**config.QueryTraceRateLimitConfig) + *cfg = &config.QueryTraceRateLimitConfig{ + SpaceMaxQPS: map[string]int{"space1": 100}, + DefaultMaxQPS: 50, + } + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background(), key: "space1"}, + want: 100, + wantErr: false, + }, + { + name: "get query max qps with default value", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), queryTraceRateLimitCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(**config.QueryTraceRateLimitConfig) + *cfg = &config.QueryTraceRateLimitConfig{ + SpaceMaxQPS: map[string]int{"space1": 100}, + DefaultMaxQPS: 50, + } + return nil + }) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background(), key: "space2"}, + want: 50, + wantErr: false, + }, + { + name: "unmarshal key failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(gomock.Any(), queryTraceRateLimitCfgKey, gomock.Any()). + Return(fmt.Errorf("unmarshal error")) + return fields{configLoader: mockLoader} + }, + args: args{ctx: context.Background(), key: "space1"}, + want: 0, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + tr := &TraceConfigCenter{ + IConfigLoader: f.configLoader, + } + got, err := tr.GetQueryMaxQPS(tt.args.ctx, tt.args.key) + assert.Equal(t, tt.wantErr, err != nil) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestNewTraceConfigCenter(t *testing.T) { + type args struct { + confP *confmocks.MockIConfigLoader + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) args + wantPanic bool + }{ + { + name: "create trace config center successfully", + fieldsGetter: func(ctrl *gomock.Controller) args { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(context.Background(), tenantTablesCfgKey, gomock.Any()). + DoAndReturn(func(ctx context.Context, key string, v interface{}, opts ...interface{}) error { + cfg := v.(**config.TenantCfg) + *cfg = &config.TenantCfg{ + DefaultIngestTenant: "test_tenant", + } + return nil + }) + return args{confP: mockLoader} + }, + wantPanic: false, + }, + { + name: "create trace config center with panic", + fieldsGetter: func(ctrl *gomock.Controller) args { + mockLoader := confmocks.NewMockIConfigLoader(ctrl) + mockLoader.EXPECT().UnmarshalKey(context.Background(), tenantTablesCfgKey, gomock.Any()). + Return(fmt.Errorf("config error")) + return args{confP: mockLoader} + }, + wantPanic: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + if tt.wantPanic { + assert.Panics(t, func() { + NewTraceConfigCenter(f.confP) + }) + } else { + got := NewTraceConfigCenter(f.confP) + assert.NotNil(t, got) + assert.IsType(t, &TraceConfigCenter{}, got) + } + }) + } +} + +// Helper function to create string pointer +func stringPtr(s string) *string { + return &s +} diff --git a/backend/modules/observability/infra/metrics/metrics.go b/backend/modules/observability/infra/metrics/metrics.go index 5091d24ee..24d72ed62 100644 --- a/backend/modules/observability/infra/metrics/metrics.go +++ b/backend/modules/observability/infra/metrics/metrics.go @@ -16,11 +16,9 @@ import ( const ( traceSpansMetricsName = "trace_spans" - getTraceSuffix = "get_trace" - listSpansSuffix = "list_spans" - searchTraceOApiSuffix = "search_trace_oapi" - listSpansOApiSuffix = "list_spans_oapi" - listTracesOApiSuffix = "list_traces_oapi" + getTraceSuffix = "get_trace" + listSpansSuffix = "list_spans" + traceOApiSuffix = "trace_oapi" throughputSuffix = ".throughput" latencySuffix = ".latency" @@ -28,15 +26,17 @@ const ( ) const ( + tagMethod = "method" tagSpaceID = "workspace_id" tagPlatformType = "platform_type" - tagSpanType = "span_type" + tagSpanType = "span_list_type" tagIsErr = "is_err" tagErrCode = "err_code" ) func traceQueryTagNames() []string { return []string{ + tagMethod, tagSpaceID, tagPlatformType, tagSpanType, @@ -102,49 +102,20 @@ func (t *TraceMetricsImpl) EmitGetTrace(workspaceId int64, start time.Time, isEr metrics.Timer(time.Since(start).Microseconds(), metrics.WithSuffix(getTraceSuffix+latencySuffix))) } -func (t *TraceMetricsImpl) EmitListSpansOapi(workspaceId int64, platformType, spanType string, spanSize int64, errorCode int, start time.Time, isError bool) { +func (t *TraceMetricsImpl) EmitTraceOapi(method string, workspaceId int64, platformType, spanListType string, spanSize int64, errorCode int, start time.Time, isError bool) { if t.spansMetrics == nil { return } t.spansMetrics.Emit( []metrics.T{ + {Name: tagMethod, Value: method}, {Name: tagSpaceID, Value: strconv.FormatInt(workspaceId, 10)}, {Name: tagIsErr, Value: strconv.FormatBool(isError)}, {Name: tagPlatformType, Value: platformType}, - {Name: tagSpanType, Value: spanType}, - {Name: tagErrCode, Value: strconv.Itoa(errorCode)}, - }, - metrics.Counter(1, metrics.WithSuffix(listSpansOApiSuffix+throughputSuffix)), - metrics.Counter(spanSize, metrics.WithSuffix(listSpansOApiSuffix+sizeSuffix)), - metrics.Timer(time.Since(start).Microseconds(), metrics.WithSuffix(listSpansOApiSuffix+latencySuffix))) -} - -func (t *TraceMetricsImpl) EmitSearchTraceOapi(workspaceId int64, platformType string, spanSize int64, errorCode int, start time.Time, isError bool) { - if t.spansMetrics == nil { - return - } - t.spansMetrics.Emit( - []metrics.T{ - {Name: tagSpaceID, Value: strconv.FormatInt(workspaceId, 10)}, - {Name: tagPlatformType, Value: platformType}, - {Name: tagIsErr, Value: strconv.FormatBool(isError)}, - {Name: tagErrCode, Value: strconv.Itoa(errorCode)}, - }, - metrics.Counter(1, metrics.WithSuffix(searchTraceOApiSuffix+throughputSuffix)), - metrics.Counter(spanSize, metrics.WithSuffix(searchTraceOApiSuffix+sizeSuffix)), - metrics.Timer(time.Since(start).Microseconds(), metrics.WithSuffix(searchTraceOApiSuffix+latencySuffix))) -} - -func (t *TraceMetricsImpl) EmitListTracesOapi(workspaceId int64, errorCode int, start time.Time, isError bool) { - if t.spansMetrics == nil { - return - } - t.spansMetrics.Emit( - []metrics.T{ - {Name: tagSpaceID, Value: strconv.FormatInt(workspaceId, 10)}, - {Name: tagIsErr, Value: strconv.FormatBool(isError)}, + {Name: tagSpanType, Value: spanListType}, {Name: tagErrCode, Value: strconv.Itoa(errorCode)}, }, - metrics.Counter(1, metrics.WithSuffix(listTracesOApiSuffix+throughputSuffix)), - metrics.Timer(time.Since(start).Microseconds(), metrics.WithSuffix(listTracesOApiSuffix+latencySuffix))) + metrics.Counter(1, metrics.WithSuffix(traceOApiSuffix+throughputSuffix)), + metrics.Counter(spanSize, metrics.WithSuffix(traceOApiSuffix+sizeSuffix)), + metrics.Timer(time.Since(start).Microseconds(), metrics.WithSuffix(traceOApiSuffix+latencySuffix))) } diff --git a/backend/modules/observability/infra/metrics/metrics_test.go b/backend/modules/observability/infra/metrics/metrics_test.go index f939139fc..0d275031d 100644 --- a/backend/modules/observability/infra/metrics/metrics_test.go +++ b/backend/modules/observability/infra/metrics/metrics_test.go @@ -194,11 +194,12 @@ func TestTraceMetricsImpl_EmitGetTrace(t *testing.T) { } } -func TestTraceMetricsImpl_EmitListSpansOapi(t *testing.T) { +func TestTraceMetricsImpl_EmitTraceOapi(t *testing.T) { type fields struct { spansMetrics infraMetrics.Metric } type args struct { + method string workspaceId int64 platformType string spanType string @@ -211,6 +212,7 @@ func TestTraceMetricsImpl_EmitListSpansOapi(t *testing.T) { name string fieldsGetter func(ctrl *gomock.Controller) fields args args + expectTags func(t *testing.T, tags []infraMetrics.T) }{ { name: "should not panic when spansMetrics is nil", @@ -219,97 +221,126 @@ func TestTraceMetricsImpl_EmitListSpansOapi(t *testing.T) { spansMetrics: nil, } }, - args: args{123, "coze", "llm", 1024, 0, time.Now(), false}, + args: args{"ListSpans", 123, "coze", "llm", 1024, 0, time.Now(), false}, }, { - name: "should emit metrics when spansMetrics is not nil", + name: "should emit metrics with correct tags for ListSpans", fieldsGetter: func(ctrl *gomock.Controller) fields { m := mocks.NewMockMetric(ctrl) - m.EXPECT().Emit(gomock.Any(), gomock.Any()).Times(1) + m.EXPECT().Emit(gomock.Any(), gomock.Any()).Do(func(tags []infraMetrics.T, values ...interface{}) { + // 验证标签 + expectedTags := map[string]string{ + tagMethod: "ListSpans", + tagSpaceID: "123", + tagIsErr: "false", + tagPlatformType: "coze", + tagSpanType: "llm", + tagErrCode: "0", + } + assert.Len(t, tags, 6) + for _, tag := range tags { + expectedValue, exists := expectedTags[tag.Name] + assert.True(t, exists, "Unexpected tag: %s", tag.Name) + assert.Equal(t, expectedValue, tag.Value, "Tag %s has wrong value", tag.Name) + } + // 验证指标值：3个指标（throughput counter, size counter, latency timer） + assert.Len(t, values, 3) + }).Times(1) return fields{ spansMetrics: m, } }, - args: args{123, "coze", "llm", 1024, 0, time.Now(), false}, + args: args{"ListSpans", 123, "coze", "llm", 1024, 0, time.Now(), false}, }, { - name: "should emit metrics with error", + name: "should emit metrics with correct tags for GetTrace", fieldsGetter: func(ctrl *gomock.Controller) fields { m := mocks.NewMockMetric(ctrl) - m.EXPECT().Emit(gomock.Any(), gomock.Any()).Times(1) + m.EXPECT().Emit(gomock.Any(), gomock.Any()).Do(func(tags []infraMetrics.T, values ...interface{}) { + expectedTags := map[string]string{ + tagMethod: "GetTrace", + tagSpaceID: "456", + tagIsErr: "false", + tagPlatformType: "dify", + tagSpanType: "workflow", + tagErrCode: "0", + } + assert.Len(t, tags, 6) + for _, tag := range tags { + expectedValue, exists := expectedTags[tag.Name] + assert.True(t, exists, "Unexpected tag: %s", tag.Name) + assert.Equal(t, expectedValue, tag.Value, "Tag %s has wrong value", tag.Name) + } + }).Times(1) return fields{ spansMetrics: m, } }, - args: args{456, "openai", "chat", 2048, 500, time.Now(), true}, + args: args{"GetTrace", 456, "dify", "workflow", 2048, 0, time.Now(), false}, }, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Cleanup(func() { - singletonTraceMetrics = nil - traceMetricsOnce = sync.Once{} - }) - ctrl := gomock.NewController(t) - defer ctrl.Finish() - fields := tt.fieldsGetter(ctrl) - tr := &TraceMetricsImpl{ - spansMetrics: fields.spansMetrics, - } - assert.NotPanics(t, func() { - tr.EmitListSpansOapi(tt.args.workspaceId, tt.args.platformType, tt.args.spanType, tt.args.spanSize, tt.args.errorCode, tt.args.start, tt.args.isError) - }) - }) - } -} - -func TestTraceMetricsImpl_EmitSearchTraceOapi(t *testing.T) { - type fields struct { - spansMetrics infraMetrics.Metric - } - type args struct { - workspaceId int64 - platformType string - spanSize int64 - errorCode int - start time.Time - isError bool - } - tests := []struct { - name string - fieldsGetter func(ctrl *gomock.Controller) fields - args args - }{ { - name: "should not panic when spansMetrics is nil", + name: "should emit metrics with error tags", fieldsGetter: func(ctrl *gomock.Controller) fields { + m := mocks.NewMockMetric(ctrl) + m.EXPECT().Emit(gomock.Any(), gomock.Any()).Do(func(tags []infraMetrics.T, values ...interface{}) { + expectedTags := map[string]string{ + tagMethod: "ListSpans", + tagSpaceID: "789", + tagIsErr: "true", + tagPlatformType: "openai", + tagSpanType: "chat", + tagErrCode: "500", + } + for _, tag := range tags { + expectedValue, exists := expectedTags[tag.Name] + assert.True(t, exists, "Unexpected tag: %s", tag.Name) + assert.Equal(t, expectedValue, tag.Value, "Tag %s has wrong value", tag.Name) + } + }).Times(1) return fields{ - spansMetrics: nil, + spansMetrics: m, } }, - args: args{123, "coze", 512, 0, time.Now(), false}, + args: args{"ListSpans", 789, "openai", "chat", 512, 500, time.Now(), true}, }, { - name: "should emit metrics when spansMetrics is not nil", + name: "should handle empty method and platform type", fieldsGetter: func(ctrl *gomock.Controller) fields { m := mocks.NewMockMetric(ctrl) - m.EXPECT().Emit(gomock.Any(), gomock.Any()).Times(1) + m.EXPECT().Emit(gomock.Any(), gomock.Any()).Do(func(tags []infraMetrics.T, values ...interface{}) { + expectedTags := map[string]string{ + tagMethod: "", + tagSpaceID: "0", + tagIsErr: "false", + tagPlatformType: "", + tagSpanType: "", + tagErrCode: "0", + } + for _, tag := range tags { + expectedValue, exists := expectedTags[tag.Name] + assert.True(t, exists, "Unexpected tag: %s", tag.Name) + assert.Equal(t, expectedValue, tag.Value, "Tag %s has wrong value", tag.Name) + } + }).Times(1) return fields{ spansMetrics: m, } }, - args: args{123, "coze", 512, 0, time.Now(), false}, + args: args{"", 0, "", "", 0, 0, time.Now(), false}, }, { - name: "should emit metrics with error", + name: "should handle negative span size", fieldsGetter: func(ctrl *gomock.Controller) fields { m := mocks.NewMockMetric(ctrl) - m.EXPECT().Emit(gomock.Any(), gomock.Any()).Times(1) + m.EXPECT().Emit(gomock.Any(), gomock.Any()).Do(func(tags []infraMetrics.T, values ...interface{}) { + // 验证负数span size也能正确处理 + assert.Len(t, values, 3) + }).Times(1) return fields{ spansMetrics: m, } }, - args: args{789, "openai", 1024, 400, time.Now(), true}, + args: args{"GetTrace", 999, "coze", "agent", -100, 0, time.Now(), false}, }, } for _, tt := range tests { @@ -325,80 +356,106 @@ func TestTraceMetricsImpl_EmitSearchTraceOapi(t *testing.T) { spansMetrics: fields.spansMetrics, } assert.NotPanics(t, func() { - tr.EmitSearchTraceOapi(tt.args.workspaceId, tt.args.platformType, tt.args.spanSize, tt.args.errorCode, tt.args.start, tt.args.isError) + tr.EmitTraceOapi(tt.args.method, tt.args.workspaceId, tt.args.platformType, tt.args.spanType, tt.args.spanSize, tt.args.errorCode, tt.args.start, tt.args.isError) }) }) } } -func TestTraceMetricsImpl_EmitListTracesOapi(t *testing.T) { - type fields struct { - spansMetrics infraMetrics.Metric +func TestTraceMetricsImpl_EmitTraceOapi_MetricValues(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockMetric := mocks.NewMockMetric(ctrl) + + // 测试指标值的正确性 + mockMetric.EXPECT().Emit(gomock.Any(), gomock.Any()).Do(func(tags []infraMetrics.T, values ...interface{}) { + // 验证有3个指标值：throughput counter, size counter, latency timer + assert.Len(t, values, 3) + + // 验证指标类型和后缀 + throughputCounter := values[0] + sizeCounter := values[1] + latencyTimer := values[2] + + assert.NotNil(t, throughputCounter) + assert.NotNil(t, sizeCounter) + assert.NotNil(t, latencyTimer) + }).Times(1) + + tr := &TraceMetricsImpl{ + spansMetrics: mockMetric, } - type args struct { - workspaceId int64 - errorCode int - start time.Time - isError bool + + start := time.Now() + tr.EmitTraceOapi("TestMethod", 12345, "test_platform", "test_span", 1024, 200, start, true) +} + +func TestTraceMetricsImpl_EmitTraceOapi_Integration(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + // 测试与domain层接口的集成 + mockMetric := mocks.NewMockMetric(ctrl) + mockMetric.EXPECT().Emit(gomock.Any(), gomock.Any()).AnyTimes() + + tr := &TraceMetricsImpl{ + spansMetrics: mockMetric, } - tests := []struct { - name string - fieldsGetter func(ctrl *gomock.Controller) fields - args args - }{ - { - name: "should not panic when spansMetrics is nil", - fieldsGetter: func(ctrl *gomock.Controller) fields { - return fields{ - spansMetrics: nil, - } - }, - args: args{123, 0, time.Now(), false}, - }, - { - name: "should emit metrics when spansMetrics is not nil", - fieldsGetter: func(ctrl *gomock.Controller) fields { - m := mocks.NewMockMetric(ctrl) - m.EXPECT().Emit(gomock.Any(), gomock.Any()).Times(1) - return fields{ - spansMetrics: m, - } - }, - args: args{123, 0, time.Now(), false}, - }, - { - name: "should emit metrics with error", - fieldsGetter: func(ctrl *gomock.Controller) fields { - m := mocks.NewMockMetric(ctrl) - m.EXPECT().Emit(gomock.Any(), gomock.Any()).Times(1) - return fields{ - spansMetrics: m, - } - }, - args: args{456, 404, time.Now(), true}, - }, + + // 验证实现了domain接口 + var domainMetrics metrics2.ITraceMetrics = tr + assert.NotNil(t, domainMetrics) + + // 测试所有接口方法 + start := time.Now() + domainMetrics.EmitListSpans(123, "llm", start, false) + domainMetrics.EmitGetTrace(456, start, true) + domainMetrics.EmitTraceOapi("TestMethod", 789, "coze", "workflow", 2048, 0, start, false) +} + +func TestTraceMetricsImpl_ConcurrentAccess(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockMetric := mocks.NewMockMetric(ctrl) + mockMetric.EXPECT().Emit(gomock.Any(), gomock.Any()).AnyTimes() + + tr := &TraceMetricsImpl{ + spansMetrics: mockMetric, } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Cleanup(func() { - singletonTraceMetrics = nil - traceMetricsOnce = sync.Once{} - }) - ctrl := gomock.NewController(t) - defer ctrl.Finish() - fields := tt.fieldsGetter(ctrl) - tr := &TraceMetricsImpl{ - spansMetrics: fields.spansMetrics, - } - assert.NotPanics(t, func() { - tr.EmitListTracesOapi(tt.args.workspaceId, tt.args.errorCode, tt.args.start, tt.args.isError) - }) - }) + + // 并发安全性测试 + concurrency := 50 + done := make(chan bool, concurrency) + + for i := 0; i < concurrency; i++ { + go func(id int) { + defer func() { done <- true }() + + start := time.Now() + workspaceId := int64(id + 1000) + + // 并发调用EmitTraceOapi方法 + tr.EmitTraceOapi("ConcurrentTest", workspaceId, "test_platform", "test_span", int64(id*10), id%2, start, id%2 == 1) + }(i) + } + + // 等待所有goroutine完成 + timeout := time.After(5 * time.Second) + for i := 0; i < concurrency; i++ { + select { + case <-done: + // 成功完成 + case <-timeout: + t.Fatal("Concurrent test timed out") + } } } func TestTraceQueryTagNames(t *testing.T) { expected := []string{ + tagMethod, tagSpaceID, tagPlatformType, tagSpanType, @@ -407,5 +464,5 @@ func TestTraceQueryTagNames(t *testing.T) { } result := traceQueryTagNames() assert.Equal(t, expected, result) - assert.Len(t, result, 5) + assert.Len(t, result, 6) } diff --git a/backend/modules/observability/infra/rpc/auth/auth.go b/backend/modules/observability/infra/rpc/auth/auth.go index d8b8c58c2..287ec809c 100644 --- a/backend/modules/observability/infra/rpc/auth/auth.go +++ b/backend/modules/observability/infra/rpc/auth/auth.go @@ -118,7 +118,7 @@ func (a *AuthProviderImpl) CheckIngestPermission(ctx context.Context, workspaceI } func (a *AuthProviderImpl) CheckQueryPermission(ctx context.Context, workspaceId, platformType string) error { - return a.CheckWorkspacePermission(ctx, rpc.AuthActionTraceRead, workspaceId) + return a.CheckWorkspacePermission(ctx, rpc.AuthActionTraceList, workspaceId) } func NewAuthProvider(cli authservice.Client) rpc.IAuthProvider { diff --git a/backend/modules/observability/infra/rpc/auth/auth_test.go b/backend/modules/observability/infra/rpc/auth/auth_test.go new file mode 100755 index 000000000..6e549d085 --- /dev/null +++ b/backend/modules/observability/infra/rpc/auth/auth_test.go @@ -0,0 +1,586 @@ +// Copyright (c) 2025 coze-dev Authors +// SPDX-License-Identifier: Apache-2.0 + +package auth + +import ( + "context" + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + + "github.com/coze-dev/coze-loop/backend/kitex_gen/base" + "github.com/coze-dev/coze-loop/backend/kitex_gen/coze/loop/foundation/auth" + authentity "github.com/coze-dev/coze-loop/backend/kitex_gen/coze/loop/foundation/domain/auth" + "github.com/coze-dev/coze-loop/backend/modules/observability/domain/component/rpc" + "github.com/coze-dev/coze-loop/backend/modules/observability/infra/rpc/auth/mocks" + "github.com/coze-dev/coze-loop/backend/pkg/lang/ptr" +) + +func TestAuthProviderImpl_CheckWorkspacePermission(t *testing.T) { + type fields struct { + cli *mocks.MockClient + } + type args struct { + ctx context.Context + action string + workspaceId string + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + wantErr bool + }{ + { + name: "check workspace permission successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, req *auth.MCheckPermissionRequest, opts ...interface{}) (*auth.MCheckPermissionResponse, error) { + return &auth.MCheckPermissionResponse{ + BaseResp: &base.BaseResp{StatusCode: 0}, + AuthRes: []*authentity.SubjectActionObjectAuthRes{ + {IsAllowed: ptr.Of(true)}, + }, + }, nil + }) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + }, + wantErr: false, + }, + { + name: "workspace id conversion failed - non-numeric string", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "invalid_id", + }, + wantErr: true, + }, + { + name: "rpc call failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("rpc error")) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + }, + wantErr: true, + }, + { + name: "response is nil", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, nil) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + }, + wantErr: true, + }, + { + name: "response status code non-zero", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&auth.MCheckPermissionResponse{ + BaseResp: &base.BaseResp{StatusCode: 500}, + }, nil) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + }, + wantErr: true, + }, + { + name: "permission denied - IsAllowed false", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&auth.MCheckPermissionResponse{ + BaseResp: &base.BaseResp{StatusCode: 0}, + AuthRes: []*authentity.SubjectActionObjectAuthRes{ + {IsAllowed: ptr.Of(false)}, + }, + }, nil) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + }, + wantErr: true, + }, + { + name: "response with nil base resp", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&auth.MCheckPermissionResponse{ + BaseResp: nil, + AuthRes: []*authentity.SubjectActionObjectAuthRes{ + {IsAllowed: ptr.Of(true)}, + }, + }, nil) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + }, + wantErr: false, + }, + { + name: "auth result is nil", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&auth.MCheckPermissionResponse{ + BaseResp: &base.BaseResp{StatusCode: 0}, + AuthRes: []*authentity.SubjectActionObjectAuthRes{ + nil, + }, + }, nil) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + }, + wantErr: false, + }, + { + name: "empty auth results", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&auth.MCheckPermissionResponse{ + BaseResp: &base.BaseResp{StatusCode: 0}, + AuthRes: []*authentity.SubjectActionObjectAuthRes{}, + }, nil) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + }, + wantErr: false, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + a := &AuthProviderImpl{ + cli: f.cli, + } + err := a.CheckWorkspacePermission(tt.args.ctx, tt.args.action, tt.args.workspaceId) + assert.Equal(t, tt.wantErr, err != nil) + }) + } +} + +func TestAuthProviderImpl_CheckViewPermission(t *testing.T) { + type fields struct { + cli *mocks.MockClient + } + type args struct { + ctx context.Context + action string + workspaceId string + viewId string + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + wantErr bool + }{ + { + name: "check view permission successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, req *auth.MCheckPermissionRequest, opts ...interface{}) (*auth.MCheckPermissionResponse, error) { + return &auth.MCheckPermissionResponse{ + BaseResp: &base.BaseResp{StatusCode: 0}, + AuthRes: []*authentity.SubjectActionObjectAuthRes{ + {IsAllowed: ptr.Of(true)}, + }, + }, nil + }) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + viewId: "view123", + }, + wantErr: false, + }, + { + name: "workspace id conversion failed - non-numeric string", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "invalid_id", + viewId: "view123", + }, + wantErr: true, + }, + { + name: "rpc call failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("rpc error")) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + viewId: "view123", + }, + wantErr: true, + }, + { + name: "response is nil", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, nil) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + viewId: "view123", + }, + wantErr: true, + }, + { + name: "response status code non-zero", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&auth.MCheckPermissionResponse{ + BaseResp: &base.BaseResp{StatusCode: 500}, + }, nil) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + viewId: "view123", + }, + wantErr: true, + }, + { + name: "permission denied - IsAllowed false", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&auth.MCheckPermissionResponse{ + BaseResp: &base.BaseResp{StatusCode: 0}, + AuthRes: []*authentity.SubjectActionObjectAuthRes{ + {IsAllowed: ptr.Of(false)}, + }, + }, nil) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + action: "read", + workspaceId: "12345", + viewId: "view123", + }, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + a := &AuthProviderImpl{ + cli: f.cli, + } + err := a.CheckViewPermission(tt.args.ctx, tt.args.action, tt.args.workspaceId, tt.args.viewId) + assert.Equal(t, tt.wantErr, err != nil) + }) + } +} + +func TestAuthProviderImpl_CheckIngestPermission(t *testing.T) { + type fields struct { + cli *mocks.MockClient + } + type args struct { + ctx context.Context + workspaceId string + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + wantErr bool + }{ + { + name: "check ingest permission successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, req *auth.MCheckPermissionRequest, opts ...interface{}) (*auth.MCheckPermissionResponse, error) { + // Verify the action is correct + assert.Equal(t, rpc.AuthActionTraceIngest, *req.Auths[0].Action) + return &auth.MCheckPermissionResponse{ + BaseResp: &base.BaseResp{StatusCode: 0}, + AuthRes: []*authentity.SubjectActionObjectAuthRes{ + {IsAllowed: ptr.Of(true)}, + }, + }, nil + }) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + workspaceId: "12345", + }, + wantErr: false, + }, + { + name: "check ingest permission failed - workspace permission check failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("rpc error")) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + workspaceId: "12345", + }, + wantErr: true, + }, + { + name: "check ingest permission failed - invalid workspace id", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + workspaceId: "invalid_id", + }, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + a := &AuthProviderImpl{ + cli: f.cli, + } + err := a.CheckIngestPermission(tt.args.ctx, tt.args.workspaceId) + assert.Equal(t, tt.wantErr, err != nil) + }) + } +} + +func TestAuthProviderImpl_CheckQueryPermission(t *testing.T) { + type fields struct { + cli *mocks.MockClient + } + type args struct { + ctx context.Context + workspaceId string + platformType string + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) fields + args args + wantErr bool + }{ + { + name: "check query permission successfully", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + DoAndReturn(func(ctx context.Context, req *auth.MCheckPermissionRequest, opts ...interface{}) (*auth.MCheckPermissionResponse, error) { + // Verify the action is correct + assert.Equal(t, rpc.AuthActionTraceList, *req.Auths[0].Action) + return &auth.MCheckPermissionResponse{ + BaseResp: &base.BaseResp{StatusCode: 0}, + AuthRes: []*authentity.SubjectActionObjectAuthRes{ + {IsAllowed: ptr.Of(true)}, + }, + }, nil + }) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + workspaceId: "12345", + platformType: "coze", + }, + wantErr: false, + }, + { + name: "check query permission failed - workspace permission check failed", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("rpc error")) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + workspaceId: "12345", + platformType: "coze", + }, + wantErr: true, + }, + { + name: "check query permission failed - invalid workspace id", + fieldsGetter: func(ctrl *gomock.Controller) fields { + mockClient := mocks.NewMockClient(ctrl) + return fields{cli: mockClient} + }, + args: args{ + ctx: context.Background(), + workspaceId: "invalid_id", + platformType: "coze", + }, + wantErr: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + a := &AuthProviderImpl{ + cli: f.cli, + } + err := a.CheckQueryPermission(tt.args.ctx, tt.args.workspaceId, tt.args.platformType) + assert.Equal(t, tt.wantErr, err != nil) + }) + } +} + +func TestNewAuthProvider(t *testing.T) { + type args struct { + cli *mocks.MockClient + } + tests := []struct { + name string + fieldsGetter func(ctrl *gomock.Controller) args + want rpc.IAuthProvider + }{ + { + name: "create new auth provider successfully", + fieldsGetter: func(ctrl *gomock.Controller) args { + mockClient := mocks.NewMockClient(ctrl) + return args{cli: mockClient} + }, + want: &AuthProviderImpl{}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + f := tt.fieldsGetter(ctrl) + got := NewAuthProvider(f.cli) + assert.NotNil(t, got) + assert.IsType(t, &AuthProviderImpl{}, got) + }) + } +} + +func TestAuthProviderImpl_Interface(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + mockClient := mocks.NewMockClient(ctrl) + + // Verify that AuthProviderImpl implements IAuthProvider interface + var _ rpc.IAuthProvider = &AuthProviderImpl{} + + provider := NewAuthProvider(mockClient) + assert.NotNil(t, provider) + assert.IsType(t, &AuthProviderImpl{}, provider) +} + +func TestAuthProviderImpl_ErrorCodes(t *testing.T) { + ctrl := gomock.NewController(t) + defer ctrl.Finish() + + mockClient := mocks.NewMockClient(ctrl) + provider := &AuthProviderImpl{cli: mockClient} + + // Test invalid workspace ID error code + err := provider.CheckWorkspacePermission(context.Background(), "read", "invalid") + assert.NotNil(t, err) + + // Check if error is wrapped with correct error code + assert.Contains(t, err.Error(), "internal error") + assert.Contains(t, err.Error(), "internal error") + // Test RPC error code + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(nil, fmt.Errorf("rpc error")) + + err = provider.CheckWorkspacePermission(context.Background(), "read", "12345") + assert.NotNil(t, err) + + // Test permission denied error code + mockClient.EXPECT().MCheckPermission(gomock.Any(), gomock.Any(), gomock.Any()). + Return(&auth.MCheckPermissionResponse{ + BaseResp: &base.BaseResp{StatusCode: 0}, + AuthRes: []*authentity.SubjectActionObjectAuthRes{ + {IsAllowed: ptr.Of(false)}, + }, + }, nil) + + err = provider.CheckWorkspacePermission(context.Background(), "read", "12345") + assert.NotNil(t, err) +} diff --git a/backend/modules/observability/infra/rpc/auth/mocks/auth_client_mock.go b/backend/modules/observability/infra/rpc/auth/mocks/auth_client_mock.go new file mode 100644 index 000000000..26aab398c --- /dev/null +++ b/backend/modules/observability/infra/rpc/auth/mocks/auth_client_mock.go @@ -0,0 +1,10 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: auth.go +// +// Generated by this command: +// +// mockgen -source=auth.go -destination=mocks/auth_client_mock.go -package=mocks +// + +// Package mocks is a generated GoMock package. +package mocks diff --git a/backend/modules/observability/infra/rpc/auth/mocks/authservice_client_mock.go b/backend/modules/observability/infra/rpc/auth/mocks/authservice_client_mock.go new file mode 100644 index 000000000..70f8f15a6 --- /dev/null +++ b/backend/modules/observability/infra/rpc/auth/mocks/authservice_client_mock.go @@ -0,0 +1,63 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: github.com/coze-dev/coze-loop/backend/kitex_gen/coze/loop/foundation/auth/authservice (interfaces: Client) +// +// Generated by this command: +// +// mockgen -destination=mocks/authservice_client_mock.go -package=mocks github.com/coze-dev/coze-loop/backend/kitex_gen/coze/loop/foundation/auth/authservice Client +// + +// Package mocks is a generated GoMock package. +package mocks + +import ( + context "context" + reflect "reflect" + + callopt "github.com/cloudwego/kitex/client/callopt" + auth "github.com/coze-dev/coze-loop/backend/kitex_gen/coze/loop/foundation/auth" + gomock "go.uber.org/mock/gomock" +) + +// MockClient is a mock of Client interface. +type MockClient struct { + ctrl *gomock.Controller + recorder *MockClientMockRecorder + isgomock struct{} +} + +// MockClientMockRecorder is the mock recorder for MockClient. +type MockClientMockRecorder struct { + mock *MockClient +} + +// NewMockClient creates a new mock instance. +func NewMockClient(ctrl *gomock.Controller) *MockClient { + mock := &MockClient{ctrl: ctrl} + mock.recorder = &MockClientMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockClient) EXPECT() *MockClientMockRecorder { + return m.recorder +} + +// MCheckPermission mocks base method. +func (m *MockClient) MCheckPermission(ctx context.Context, request *auth.MCheckPermissionRequest, callOptions ...callopt.Option) (*auth.MCheckPermissionResponse, error) { + m.ctrl.T.Helper() + varargs := []any{ctx, request} + for _, a := range callOptions { + varargs = append(varargs, a) + } + ret := m.ctrl.Call(m, "MCheckPermission", varargs...) + ret0, _ := ret[0].(*auth.MCheckPermissionResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// MCheckPermission indicates an expected call of MCheckPermission. +func (mr *MockClientMockRecorder) MCheckPermission(ctx, request any, callOptions ...any) *gomock.Call { + mr.mock.ctrl.T.Helper() + varargs := append([]any{ctx, request}, callOptions...) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MCheckPermission", reflect.TypeOf((*MockClient)(nil).MCheckPermission), varargs...) +} diff --git a/backend/modules/observability/infra/workspace/workspace.go b/backend/modules/observability/infra/workspace/workspace.go index fcf35a0f4..c85eed044 100644 --- a/backend/modules/observability/infra/workspace/workspace.go +++ b/backend/modules/observability/infra/workspace/workspace.go @@ -5,6 +5,7 @@ package workspace import ( "context" + "strconv" "github.com/coze-dev/coze-loop/backend/kitex_gen/coze/loop/observability/domain/span" "github.com/coze-dev/coze-loop/backend/modules/observability/domain/component/workspace" @@ -23,6 +24,6 @@ func (t *WorkspaceProviderImpl) GetIngestWorkSpaceID(ctx context.Context, spans return spans[0].WorkspaceID } -func (t *WorkspaceProviderImpl) GetQueryWorkSpaceID(ctx context.Context, requestWorkspaceID int64) int64 { - return requestWorkspaceID +func (t *WorkspaceProviderImpl) GetThirdPartyQueryWorkSpaceID(ctx context.Context, requestWorkspaceID int64) string { + return strconv.FormatInt(requestWorkspaceID, 10) } diff --git a/backend/modules/observability/infra/workspace/workspace_test.go b/backend/modules/observability/infra/workspace/workspace_test.go new file mode 100755 index 000000000..f797d2f05 --- /dev/null +++ b/backend/modules/observability/infra/workspace/workspace_test.go @@ -0,0 +1,198 @@ +// Copyright (c) 2025 coze-dev Authors +// SPDX-License-Identifier: Apache-2.0 + +package workspace + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/coze-dev/coze-loop/backend/kitex_gen/coze/loop/observability/domain/span" + "github.com/coze-dev/coze-loop/backend/modules/observability/domain/component/workspace" +) + +func TestWorkspaceProviderImpl_GetIngestWorkSpaceID(t *testing.T) { + type args struct { + ctx context.Context + spans []*span.InputSpan + } + tests := []struct { + name string + args args + want string + }{ + { + name: "empty spans array, return empty string", + args: args{ + ctx: context.Background(), + spans: []*span.InputSpan{}, + }, + want: "", + }, + { + name: "nil spans, return empty string", + args: args{ + ctx: context.Background(), + spans: nil, + }, + want: "", + }, + { + name: "normal spans array, return first span workspace id", + args: args{ + ctx: context.Background(), + spans: []*span.InputSpan{ + {WorkspaceID: "workspace1"}, + {WorkspaceID: "workspace2"}, + }, + }, + want: "workspace1", + }, + { + name: "first span has empty workspace id", + args: args{ + ctx: context.Background(), + spans: []*span.InputSpan{ + {WorkspaceID: ""}, + {WorkspaceID: "workspace2"}, + }, + }, + want: "", + }, + { + name: "single span with workspace id", + args: args{ + ctx: context.Background(), + spans: []*span.InputSpan{ + {WorkspaceID: "single_workspace"}, + }, + }, + want: "single_workspace", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + w := &WorkspaceProviderImpl{} + got := w.GetIngestWorkSpaceID(tt.args.ctx, tt.args.spans) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestWorkspaceProviderImpl_GetQueryWorkSpaceID(t *testing.T) { + type args struct { + ctx context.Context + requestWorkspaceID int64 + } + tests := []struct { + name string + args args + want string + }{ + { + name: "positive workspace id conversion", + args: args{ + ctx: context.Background(), + requestWorkspaceID: 12345, + }, + want: "12345", + }, + { + name: "negative workspace id conversion", + args: args{ + ctx: context.Background(), + requestWorkspaceID: -1, + }, + want: "-1", + }, + { + name: "zero workspace id conversion", + args: args{ + ctx: context.Background(), + requestWorkspaceID: 0, + }, + want: "0", + }, + { + name: "large positive workspace id conversion", + args: args{ + ctx: context.Background(), + requestWorkspaceID: 9223372036854775807, // max int64 + }, + want: "9223372036854775807", + }, + { + name: "large negative workspace id conversion", + args: args{ + ctx: context.Background(), + requestWorkspaceID: -9223372036854775808, // min int64 + }, + want: "-9223372036854775808", + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + w := &WorkspaceProviderImpl{} + got := w.GetThirdPartyQueryWorkSpaceID(tt.args.ctx, tt.args.requestWorkspaceID) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestNewWorkspaceProvider(t *testing.T) { + tests := []struct { + name string + want workspace.IWorkSpaceProvider + }{ + { + name: "create new workspace provider successfully", + want: &WorkspaceProviderImpl{}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := NewWorkspaceProvider() + assert.NotNil(t, got) + assert.IsType(t, &WorkspaceProviderImpl{}, got) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestWorkspaceProviderImpl_Interface(t *testing.T) { + // Verify that WorkspaceProviderImpl implements IWorkSpaceProvider interface + var _ workspace.IWorkSpaceProvider = &WorkspaceProviderImpl{} + + provider := NewWorkspaceProvider() + assert.NotNil(t, provider) + + // Test interface methods exist and are callable + workspaceID := provider.GetThirdPartyQueryWorkSpaceID(context.Background(), 123) + assert.Equal(t, "123", workspaceID) + + spans := []*span.InputSpan{{WorkspaceID: "test"}} + ingestID := provider.GetIngestWorkSpaceID(context.Background(), spans) + assert.Equal(t, "test", ingestID) +} + +func TestWorkspaceProviderImpl_EdgeCases(t *testing.T) { + provider := &WorkspaceProviderImpl{} + ctx := context.Background() + + // Test with nil context (should still work) + got := provider.GetThirdPartyQueryWorkSpaceID(ctx, 456) + assert.Equal(t, "456", got) + + // Test with nil context for GetIngestWorkSpaceID + spans := []*span.InputSpan{{WorkspaceID: "test_nil_ctx"}} + ingestID := provider.GetIngestWorkSpaceID(ctx, spans) + assert.Equal(t, "test_nil_ctx", ingestID) + + // Test with nil spans element - this would cause panic in the actual code + // So we test the actual behavior which is to return empty string for empty array + emptySpans := []*span.InputSpan{} + emptyID := provider.GetIngestWorkSpaceID(ctx, emptySpans) + assert.Equal(t, "", emptyID) +}