milvus/internal/util/grpcclient/local_grpc_client_test.go
Zhen Ye b9a10a2f68
enhance: remove the rpc layer of coordinator when enabling standalone or mixcoord (#37815)
issue: #37764

- add a local client to call local server directly for
querycoord/rootcoord/datacoord.
- enable local client if milvus is running mixcoord or standalone mode.

---------

Signed-off-by: chyezh <chyezh@outlook.com>
2024-11-23 21:36:33 +08:00

47 lines
1.2 KiB
Go

package grpcclient
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
)
type mockRootCoordServer struct {
t *testing.T
*rootcoordpb.UnimplementedRootCoordServer
}
func (s *mockRootCoordServer) AllocID(ctx context.Context, req *rootcoordpb.AllocIDRequest) (*rootcoordpb.AllocIDResponse, error) {
assert.NotNil(s.t, req)
assert.Equal(s.t, uint32(100), req.Count)
return &rootcoordpb.AllocIDResponse{
ID: 1,
Count: 2,
}, nil
}
func TestLocalGRPCClient(t *testing.T) {
localClient := NewLocalGRPCClient(
&rootcoordpb.RootCoord_ServiceDesc,
&mockRootCoordServer{
t: t,
UnimplementedRootCoordServer: &rootcoordpb.UnimplementedRootCoordServer{},
},
rootcoordpb.NewRootCoordClient,
)
result, err := localClient.AllocTimestamp(context.Background(), &rootcoordpb.AllocTimestampRequest{})
assert.Error(t, err)
assert.Nil(t, result)
result2, err := localClient.AllocID(context.Background(), &rootcoordpb.AllocIDRequest{
Count: 100,
})
assert.NoError(t, err)
assert.NotNil(t, result2)
assert.Equal(t, int64(1), result2.ID)
assert.Equal(t, uint32(2), result2.Count)
}