mirror of
https://gitee.com/milvus-io/milvus.git
synced 2025-12-06 17:18:35 +08:00
511a04a6a5
Related to #45105 This commit refactors the test MilvusClient wrapper to leverage Go's embedding pattern and improves test organization with subtests. **File**: `tests/go_client/base/milvus_client.go` - **Use `typeutil.NewSet` for rate limiting**: Replace map-based `rateLogMethods` with `typeutil.NewSet` for cleaner and more efficient membership checking - **Embed `*client.Client` directly**: Change `MilvusClient` structure from wrapping the client as a field to embedding it directly - **Remove ~380 lines of wrapper methods**: All wrapper methods (database, collection, partition, index, read/write, RBAC, etc.) are now unnecessary thanks to Go's embedding feature, which automatically promotes embedded methods to the outer type - **Simplify initialization**: Update `NewMilvusClient` and `Close` to use embedded client directly - **Fix typo**: Correct comment "Ike the actual method" → "Invoke the actual method" **File**: `tests/go_client/testcases/search_test.go` - **Wrap assertions in subtests**: Each search expression test is now wrapped in `t.Run()` with descriptive names - **Dynamic subtest naming**: Format: `expr={expression}_dynamic-{true/false}` for clear test identification Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
84 lines
2.6 KiB
Go
84 lines
2.6 KiB
Go
package base
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"strings"
|
|
"time"
|
|
|
|
"go.uber.org/zap"
|
|
"google.golang.org/grpc"
|
|
|
|
client "github.com/milvus-io/milvus/client/v2/milvusclient"
|
|
"github.com/milvus-io/milvus/pkg/v2/log"
|
|
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
|
|
)
|
|
|
|
func LoggingUnaryInterceptor() grpc.UnaryClientInterceptor {
|
|
// Limit debug logging for these methods
|
|
ratedLogMethods := typeutil.NewSet("GetFlushState", "GetLoadingProgress", "DescribeIndex")
|
|
|
|
logWithRateLimit := func(methodShortName string, logFunc func(msg string, fields ...zap.Field),
|
|
logRateFunc func(cost float64, msg string, fields ...zap.Field) bool,
|
|
msg string, fields ...zap.Field,
|
|
) {
|
|
if ratedLogMethods.Contain(methodShortName) {
|
|
logRateFunc(10, msg, fields...)
|
|
} else {
|
|
logFunc(msg, fields...)
|
|
}
|
|
}
|
|
|
|
return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
|
|
const maxLogLength = 300
|
|
_method := strings.Split(method, "/")
|
|
_methodShortName := _method[len(_method)-1]
|
|
|
|
// Marshal request
|
|
marshalWithFallback := func(v interface{}, fallbackMsg string) string {
|
|
dataJSON, err := json.Marshal(v)
|
|
if err != nil {
|
|
log.Error("Failed to marshal", zap.Error(err))
|
|
return fallbackMsg
|
|
}
|
|
dataStr := string(dataJSON)
|
|
if len(dataStr) > maxLogLength {
|
|
return dataStr[:maxLogLength] + "......"
|
|
}
|
|
return dataStr
|
|
}
|
|
|
|
reqStr := marshalWithFallback(req, "could not marshal request")
|
|
logWithRateLimit(_methodShortName, log.Info, log.RatedInfo, "Request", zap.String("method", _methodShortName), zap.String("reqs", reqStr))
|
|
|
|
// Invoke the actual method
|
|
start := time.Now()
|
|
errResp := invoker(ctx, method, req, reply, cc, opts...)
|
|
cost := time.Since(start)
|
|
|
|
// Marshal response
|
|
respStr := marshalWithFallback(reply, "could not marshal response")
|
|
logWithRateLimit(_methodShortName, log.Info, log.RatedInfo, "Response", zap.String("method", _methodShortName), zap.String("resp", respStr))
|
|
logWithRateLimit(_methodShortName, log.Debug, log.RatedDebug, "Cost", zap.String("method", _methodShortName), zap.Duration("cost", cost))
|
|
|
|
return errResp
|
|
}
|
|
}
|
|
|
|
type MilvusClient struct {
|
|
*client.Client
|
|
}
|
|
|
|
func NewMilvusClient(ctx context.Context, cfg *client.ClientConfig) (*MilvusClient, error) {
|
|
cfg.DialOptions = append(cfg.DialOptions, grpc.WithUnaryInterceptor(LoggingUnaryInterceptor()))
|
|
mClient, err := client.New(ctx, cfg)
|
|
return &MilvusClient{
|
|
Client: mClient,
|
|
}, err
|
|
}
|
|
|
|
func (mc *MilvusClient) Close(ctx context.Context) error {
|
|
err := mc.Client.Close(ctx)
|
|
return err
|
|
}
|