Files
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

500 lines
14 KiB
Go

//go:build test
package mlog
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"go.opentelemetry.io/otel/trace"
"go.uber.org/zap/zapcore"
"google.golang.org/grpc"
"google.golang.org/grpc/metadata"
)
func Test_extractPropagatedFromMetadata(t *testing.T) {
// Type-encoded keys: "mlog-s-" for string, "mlog-i-" for int64
md := metadata.New(map[string]string{
metadataPrefixString + "collectionname": "my_collection",
metadataPrefixInt64 + "collectionid": "12345",
})
ctx := metadata.NewIncomingContext(context.Background(), md)
ctx = extractPropagated(ctx)
props := GetPropagated(ctx)
assert.Len(t, props, 2)
assert.Equal(t, "my_collection", props[keyCollectionName])
assert.Equal(t, "12345", props[keyCollectionID])
// Verify int64 field preserves type
fields := FieldsFromContext(ctx)
for _, f := range fields {
if f.Key == keyCollectionID {
assert.Equal(t, zapcore.Int64Type, f.Type)
assert.Equal(t, int64(12345), f.Integer)
}
}
}
func Test_extractPropagatedIgnoresNonPrefixedKeys(t *testing.T) {
md := metadata.New(map[string]string{
metadataPrefixString + "propagated": "value",
"other-key": "other-value",
})
ctx := metadata.NewIncomingContext(context.Background(), md)
ctx = extractPropagated(ctx)
props := GetPropagated(ctx)
assert.Len(t, props, 1)
assert.Equal(t, "value", props["propagated"])
}
func Test_extractPropagatedLegacyFormat(t *testing.T) {
// Legacy format without type prefix — treated as string
md := metadata.New(map[string]string{
MetadataPrefix + "oldkey": "oldvalue",
})
ctx := metadata.NewIncomingContext(context.Background(), md)
ctx = extractPropagated(ctx)
props := GetPropagated(ctx)
assert.Len(t, props, 1)
assert.Equal(t, "oldvalue", props["oldkey"])
}
func Test_extractPropagatedNoMetadata(t *testing.T) {
ctx := context.Background()
ctx = extractPropagated(ctx)
props := GetPropagated(ctx)
assert.Nil(t, props)
}
func Test_extractPropagatedEmptyMetadata(t *testing.T) {
md := metadata.New(map[string]string{})
ctx := metadata.NewIncomingContext(context.Background(), md)
ctx = extractPropagated(ctx)
props := GetPropagated(ctx)
assert.Nil(t, props)
}
func Test_injectPropagatedToMetadata(t *testing.T) {
ctx := context.Background()
ctx = WithFields(ctx,
propagatedStringField(keyCollectionName, "my_collection"),
propagatedInt64Field(keyCollectionID, 12345),
)
ctx = injectPropagated(ctx)
md, ok := metadata.FromOutgoingContext(ctx)
assert.True(t, ok)
assert.Equal(t, []string{"my_collection"}, md.Get(metadataPrefixString+"collectionname"))
assert.Equal(t, []string{"12345"}, md.Get(metadataPrefixInt64+"collectionid"))
}
func Test_injectPropagatedNoPropagatedFields(t *testing.T) {
ctx := context.Background()
ctx = injectPropagated(ctx)
_, ok := metadata.FromOutgoingContext(ctx)
assert.False(t, ok)
}
func Test_injectPropagatedAppendsToExistingMetadata(t *testing.T) {
ctx := context.Background()
ctx = metadata.AppendToOutgoingContext(ctx, "existing-key", "existing-value")
ctx = WithFields(ctx, propagatedStringField("new-key", "new-value"))
ctx = injectPropagated(ctx)
md, ok := metadata.FromOutgoingContext(ctx)
assert.True(t, ok)
assert.Equal(t, []string{"existing-value"}, md.Get("existing-key"))
assert.Equal(t, []string{"new-value"}, md.Get(metadataPrefixString+"new-key"))
}
func TestRoundTripPropagation(t *testing.T) {
// Simulate client side: create context with propagated fields
clientCtx := context.Background()
clientCtx = WithFields(clientCtx,
propagatedStringField(keyCollectionName, "test_collection"),
propagatedInt64Field(keyCollectionID, 99999),
)
// Client injects into outgoing metadata
clientCtx = injectPropagated(clientCtx)
// Simulate network transfer: copy outgoing to incoming
outMd, _ := metadata.FromOutgoingContext(clientCtx)
serverCtx := metadata.NewIncomingContext(context.Background(), outMd)
// Server extracts from incoming metadata
serverCtx = extractPropagated(serverCtx)
// Verify propagated fields are preserved with correct types
props := GetPropagated(serverCtx)
assert.Len(t, props, 2)
assert.Equal(t, "test_collection", props[keyCollectionName])
assert.Equal(t, "99999", props[keyCollectionID])
// Verify int64 type is preserved across the round trip
for _, f := range FieldsFromContext(serverCtx) {
if f.Key == keyCollectionID {
assert.Equal(t, zapcore.Int64Type, f.Type)
assert.Equal(t, int64(99999), f.Integer)
}
if f.Key == keyCollectionName {
assert.Equal(t, zapcore.StringType, f.Type)
assert.Equal(t, "test_collection", f.String)
}
}
}
func TestMetadataPrefix(t *testing.T) {
assert.Equal(t, "mlog-", MetadataPrefix)
}
func TestUnaryServerInterceptorExtractsFields(t *testing.T) {
interceptor := UnaryServerInterceptor("test-module")
md := metadata.New(map[string]string{
metadataPrefixString + "key": "value",
})
ctx := metadata.NewIncomingContext(context.Background(), md)
var handlerCtx context.Context
handler := func(ctx context.Context, req any) (any, error) {
handlerCtx = ctx
return "response", nil
}
resp, err := interceptor(ctx, nil, nil, handler)
assert.NoError(t, err)
assert.Equal(t, "response", resp)
props := GetPropagated(handlerCtx)
assert.Equal(t, "value", props["key"])
}
func TestUnaryClientInterceptorInjectsFields(t *testing.T) {
interceptor := UnaryClientInterceptor()
ctx := context.Background()
ctx = WithFields(ctx, propagatedStringField("key", "value"))
var invokerCtx context.Context
invoker := func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, opts ...grpc.CallOption) error {
invokerCtx = ctx
return nil
}
err := interceptor(ctx, "/test", nil, nil, nil, invoker)
assert.NoError(t, err)
md, ok := metadata.FromOutgoingContext(invokerCtx)
assert.True(t, ok)
assert.Equal(t, []string{"value"}, md.Get(metadataPrefixString+"key"))
}
// mockServerStream is a minimal mock for grpc.ServerStream
type mockServerStream struct {
grpc.ServerStream
ctx context.Context
}
func (m *mockServerStream) Context() context.Context {
return m.ctx
}
func TestStreamServerInterceptorExtractsFields(t *testing.T) {
interceptor := StreamServerInterceptor("test-module")
md := metadata.New(map[string]string{
metadataPrefixString + "key": "value",
})
ctx := metadata.NewIncomingContext(context.Background(), md)
mockStream := &mockServerStream{ctx: ctx}
var handlerStream grpc.ServerStream
handler := func(srv any, stream grpc.ServerStream) error {
handlerStream = stream
return nil
}
err := interceptor(nil, mockStream, nil, handler)
assert.NoError(t, err)
// Verify the wrapped stream has the extracted context
props := GetPropagated(handlerStream.Context())
assert.Equal(t, "value", props["key"])
}
func TestStreamServerInterceptorWrappedStreamContext(t *testing.T) {
interceptor := StreamServerInterceptor("test-module")
md := metadata.New(map[string]string{
metadataPrefixString + "a": "1",
metadataPrefixString + "b": "2",
})
ctx := metadata.NewIncomingContext(context.Background(), md)
mockStream := &mockServerStream{ctx: ctx}
var handlerStream grpc.ServerStream
handler := func(srv any, stream grpc.ServerStream) error {
handlerStream = stream
return nil
}
err := interceptor(nil, mockStream, nil, handler)
assert.NoError(t, err)
// The handler receives a wrappedStream with the updated context
props := GetPropagated(handlerStream.Context())
assert.Len(t, props, 2)
assert.Equal(t, "1", props["a"])
assert.Equal(t, "2", props["b"])
}
func TestStreamClientInterceptorInjectsFields(t *testing.T) {
interceptor := StreamClientInterceptor()
ctx := context.Background()
ctx = WithFields(ctx, propagatedStringField("key", "value"))
var streamerCtx context.Context
streamer := func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) {
streamerCtx = ctx
return nil, nil
}
_, err := interceptor(ctx, nil, nil, "/test", streamer)
assert.NoError(t, err)
md, ok := metadata.FromOutgoingContext(streamerCtx)
assert.True(t, ok)
assert.Equal(t, []string{"value"}, md.Get(metadataPrefixString+"key"))
}
func TestRegularFieldsNotPropagated(t *testing.T) {
ctx := context.Background()
ctx = WithFields(ctx, String("regular", "value"))
ctx = injectPropagated(ctx)
// Regular fields should not be injected into metadata
_, ok := metadata.FromOutgoingContext(ctx)
assert.False(t, ok)
}
func Test_extractPropagatedWithTraceContext(t *testing.T) {
traceID, _ := trace.TraceIDFromHex("0102030405060708090a0b0c0d0e0f10")
spanID, _ := trace.SpanIDFromHex("0102030405060708")
spanCtx := trace.NewSpanContext(trace.SpanContextConfig{
TraceID: traceID,
SpanID: spanID,
})
ctx := trace.ContextWithSpanContext(context.Background(), spanCtx)
ctx = extractPropagated(ctx)
// TraceID and SpanID are appended dynamically at log time, not cached as context fields.
assert.Nil(t, FieldsFromContext(ctx))
}
func Test_extractPropagatedWithTraceContextAndMetadata(t *testing.T) {
traceID, _ := trace.TraceIDFromHex("0102030405060708090a0b0c0d0e0f10")
spanID, _ := trace.SpanIDFromHex("0102030405060708")
spanCtx := trace.NewSpanContext(trace.SpanContextConfig{
TraceID: traceID,
SpanID: spanID,
})
ctx := trace.ContextWithSpanContext(context.Background(), spanCtx)
// Add gRPC metadata
md := metadata.New(map[string]string{
metadataPrefixString + "collectionname": "my_collection",
})
ctx = metadata.NewIncomingContext(ctx, md)
ctx = extractPropagated(ctx)
// Verify only mlog metadata fields are cached on context.
fields := FieldsFromContext(ctx)
assert.Len(t, fields, 1)
fieldMap := make(map[string]string)
for _, f := range fields {
fieldMap[f.Key] = f.String
}
assert.NotContains(t, fieldMap, keyTraceID)
assert.NotContains(t, fieldMap, keySpanID)
// Propagated field should be accessible via GetPropagated
props := GetPropagated(ctx)
assert.Equal(t, "my_collection", props[keyCollectionName])
}
func Test_extractPropagatedNoTraceContext(t *testing.T) {
ctx := context.Background()
ctx = extractPropagated(ctx)
// No fields should be added without trace context or metadata
fields := FieldsFromContext(ctx)
assert.Nil(t, fields)
}
func Test_extractPropagatedWithExtraFields(t *testing.T) {
md := metadata.New(map[string]string{
metadataPrefixString + "key": "value",
})
ctx := metadata.NewIncomingContext(context.Background(), md)
// Pass extra fields to extractPropagated
ctx = extractPropagated(ctx,
String(keyModule, "proxy"),
Int64("custom", 123),
)
// Verify all fields are present in a single WithFields call
fields := FieldsFromContext(ctx)
assert.Len(t, fields, 3) // key, module, custom
fieldMap := make(map[string]any)
for _, f := range fields {
if f.String != "" {
fieldMap[f.Key] = f.String
} else {
fieldMap[f.Key] = f.Integer
}
}
assert.Equal(t, "proxy", fieldMap[keyModule])
assert.Equal(t, int64(123), fieldMap["custom"])
props := GetPropagated(ctx)
assert.Equal(t, "value", props["key"])
}
func Test_extractPropagatedOnlyExtraFields(t *testing.T) {
ctx := context.Background()
// Only extra fields, no metadata or trace context
ctx = extractPropagated(ctx, String(keyModule, "datanode"))
fields := FieldsFromContext(ctx)
assert.Len(t, fields, 1)
assert.Equal(t, "datanode", fields[0].String)
}
func TestUnaryServerInterceptorWithModule(t *testing.T) {
interceptor := UnaryServerInterceptor("proxy")
ctx := context.Background()
var handlerCtx context.Context
handler := func(ctx context.Context, req any) (any, error) {
handlerCtx = ctx
return "response", nil
}
resp, err := interceptor(ctx, nil, nil, handler)
assert.NoError(t, err)
assert.Equal(t, "response", resp)
// Verify module field is added
fields := FieldsFromContext(handlerCtx)
fieldMap := make(map[string]string)
for _, f := range fields {
fieldMap[f.Key] = f.String
}
assert.Equal(t, "proxy", fieldMap[keyModule])
}
func TestUnaryServerInterceptorModuleWithMetadata(t *testing.T) {
interceptor := UnaryServerInterceptor("querynode")
md := metadata.New(map[string]string{
metadataPrefixString + "collectionname": "my_collection",
})
ctx := metadata.NewIncomingContext(context.Background(), md)
var handlerCtx context.Context
handler := func(ctx context.Context, req any) (any, error) {
handlerCtx = ctx
return "response", nil
}
_, err := interceptor(ctx, nil, nil, handler)
assert.NoError(t, err)
// Verify both module and propagated fields are present
fields := FieldsFromContext(handlerCtx)
fieldMap := make(map[string]string)
for _, f := range fields {
fieldMap[f.Key] = f.String
}
assert.Equal(t, "querynode", fieldMap[keyModule])
props := GetPropagated(handlerCtx)
assert.Equal(t, "my_collection", props[keyCollectionName])
}
func TestStreamServerInterceptorWithModule(t *testing.T) {
interceptor := StreamServerInterceptor("datanode")
ctx := context.Background()
mockStream := &mockServerStream{ctx: ctx}
var handlerStream grpc.ServerStream
handler := func(srv any, stream grpc.ServerStream) error {
handlerStream = stream
return nil
}
err := interceptor(nil, mockStream, nil, handler)
assert.NoError(t, err)
// Verify module field is added
fields := FieldsFromContext(handlerStream.Context())
fieldMap := make(map[string]string)
for _, f := range fields {
fieldMap[f.Key] = f.String
}
assert.Equal(t, "datanode", fieldMap[keyModule])
}
func TestStreamServerInterceptorModuleWithMetadata(t *testing.T) {
interceptor := StreamServerInterceptor("streamingnode")
md := metadata.New(map[string]string{
metadataPrefixString + "key": "value",
})
ctx := metadata.NewIncomingContext(context.Background(), md)
mockStream := &mockServerStream{ctx: ctx}
var handlerStream grpc.ServerStream
handler := func(srv any, stream grpc.ServerStream) error {
handlerStream = stream
return nil
}
err := interceptor(nil, mockStream, nil, handler)
assert.NoError(t, err)
// Verify both module and propagated fields are present
fields := FieldsFromContext(handlerStream.Context())
fieldMap := make(map[string]string)
for _, f := range fields {
fieldMap[f.Key] = f.String
}
assert.Equal(t, "streamingnode", fieldMap[keyModule])
props := GetPropagated(handlerStream.Context())
assert.Equal(t, "value", props["key"])
}