milvus/internal/util/streamingutil/status/client_stream_wrapper_test.go
congqixia cb7f2fa6fd
enhance: Use v2 package name for pkg module (#39990)
Related to #39095

https://go.dev/doc/modules/version-numbers

Update pkg version according to golang dep version convention

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
2025-02-22 23:15:58 +08:00

34 lines
1.1 KiB
Go

package status
import (
"context"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus/internal/mocks/google.golang.org/mock_grpc"
"github.com/milvus-io/milvus/pkg/v2/proto/streamingpb"
)
func TestClientStreamWrapper(t *testing.T) {
s := mock_grpc.NewMockClientStream(t)
s.EXPECT().SendMsg(mock.Anything).Return(NewGRPCStatusFromStreamingError(NewOnShutdownError("test")).Err())
s.EXPECT().RecvMsg(mock.Anything).Return(NewGRPCStatusFromStreamingError(NewOnShutdownError("test")).Err())
w := NewClientStreamWrapper("method", s)
err := w.SendMsg(context.Background())
assert.NotNil(t, err)
streamingErr := AsStreamingError(err)
assert.Equal(t, streamingpb.StreamingCode_STREAMING_CODE_ON_SHUTDOWN, streamingErr.Code)
assert.Contains(t, streamingErr.Cause, "test")
err = w.RecvMsg(context.Background())
assert.NotNil(t, err)
streamingErr = AsStreamingError(err)
assert.Equal(t, streamingpb.StreamingCode_STREAMING_CODE_ON_SHUTDOWN, streamingErr.Code)
assert.Contains(t, streamingErr.Cause, "test")
assert.Nil(t, NewClientStreamWrapper("method", nil))
}