Files
milvus/internal/proxy/proxy_rpc_test.go
Zhen YeandGitHub b07c62d53c enhance: replace log package with context-aware mlog (#50094)
issue: #35917

- mlog package: move logger initialization, zap core, async buffered
writes, field helpers, and scoped logger binding into pkg/mlog while
removing pkg/log.
- logging callsites: migrate Milvus logging usage to context-aware mlog
APIs and simplify redundant With chains across components, utilities,
tests, and tools.
- trace propagation: replace logutil trace interceptors with mlog/tracer
integration and add client_request_id fallback propagation for server
stats handlers.

---------

Signed-off-by: chyezh <chyezh@outlook.com>
2026-06-24 00:58:28 +08:00

74 lines
2.1 KiB
Go

package proxy
import (
"context"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
grpcproxyclient "github.com/milvus-io/milvus/internal/distributed/proxy/client"
"github.com/milvus-io/milvus/internal/util/dependency"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
"github.com/milvus-io/milvus/pkg/v3/proto/proxypb"
"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
func TestProxyRpcLimit(t *testing.T) {
var err error
path := t.TempDir() + "/rocksmq"
t.Setenv("ROCKSMQ_PATH", path)
ctx := GetContext(context.Background(), "root:123456")
localMsg := true
factory := dependency.NewDefaultFactory(localMsg)
bt := paramtable.NewBaseTable(paramtable.SkipRemote(true))
base := &paramtable.ComponentParam{}
base.Init(bt)
var p paramtable.GrpcServerConfig
assert.NoError(t, err)
p.Init(typeutil.ProxyRole, bt)
base.Save("proxy.grpc.serverMaxRecvSize", "1")
assert.Equal(t, p.ServerMaxRecvSize.GetAsInt(), 1)
mlog.Info(context.TODO(), "Initialize parameter table of Proxy")
proxy, err := NewProxy(ctx, factory)
assert.NoError(t, err)
assert.NotNil(t, proxy)
testServer := newProxyTestServer(proxy)
go testServer.startGrpc(ctx, &p)
assert.NoError(t, testServer.waitForGrpcReady())
testServer.SetAddress(testServer.lisAddr)
defer testServer.grpcServer.Stop()
client, err := grpcproxyclient.NewClient(ctx, testServer.lisAddr, 1)
assert.NoError(t, err)
proxy.UpdateStateCode(commonpb.StateCode_Healthy)
rates := make([]*internalpb.Rate, 0)
req := &proxypb.SetRatesRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgID(int64(0)),
commonpbutil.WithTimeStamp(0),
),
Rates: []*proxypb.CollectionRate{
{
Collection: 1,
Rates: rates,
},
},
}
_, err = client.SetRates(ctx, req)
// should be limited because of the rpc limit
assert.Error(t, err)
assert.True(t, strings.Contains(err.Error(), "ResourceExhausted"))
}