package conn

import (
	"context"
	"errors"
	"fmt"
	"io"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Query"
	"go.uber.org/mock/gomock"
	grpcCodes "google.golang.org/grpc/codes"
	"google.golang.org/grpc/metadata"
	grpcStatus "google.golang.org/grpc/status"

	"github.com/ydb-platform/ydb-go-sdk/v3/internal/backoff"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/mock"
	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
	"github.com/ydb-platform/ydb-go-sdk/v3/retry"
	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
)

func TestGrpcClientStream_Header(t *testing.T) {
	t.Run("ReturnsHeaderFromUnderlyingStream", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		expectedMD := metadata.MD{"key": []string{"value"}}
		mockStream.EXPECT().Header().Return(expectedMD, nil)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
		}

		md, err := s.Header()
		require.NoError(t, err)
		require.Equal(t, expectedMD, md)
	})

	t.Run("ReturnsErrorFromUnderlyingStream", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		expectedErr := fmt.Errorf("header error")
		mockStream.EXPECT().Header().Return(metadata.MD{}, expectedErr)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
		}

		md, err := s.Header()
		require.Error(t, err)
		require.Empty(t, md)
	})
}

func TestGrpcClientStream_Trailer(t *testing.T) {
	t.Run("ReturnsTrailerFromUnderlyingStream", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		expectedMD := metadata.MD{"trailer-key": []string{"trailer-value"}}
		mockStream.EXPECT().Trailer().Return(expectedMD)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
		}

		md := s.Trailer()
		require.Equal(t, expectedMD, md)
	})
}

func TestGrpcClientStream_Context(t *testing.T) {
	t.Run("ReturnsContextFromUnderlyingStream", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		expectedCtx := context.WithValue(t.Context(), "key", "value") //nolint:revive,staticcheck
		mockStream.EXPECT().Context().Return(expectedCtx)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
		}

		ctx := s.Context()
		require.Equal(t, expectedCtx, ctx)
	})
}

func TestGrpcClientStream_CloseSend(t *testing.T) {
	t.Run("SuccessfulClose", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		mockStream.EXPECT().CloseSend().Return(nil)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   true,
			traceID:    "test-trace-id",
		}

		err := s.CloseSend()
		require.NoError(t, err)
	})

	t.Run("ContextError", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		mockStream.EXPECT().CloseSend().Return(context.Canceled)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   true,
			traceID:    "test-trace-id",
		}

		err := s.CloseSend()
		require.Error(t, err)
		require.True(t, xerrors.IsContextError(err))
	})

	t.Run("StreamContextDoneReturnsNonTransportError", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		// Use a non-gRPC-status error that gRPC may return on stream termination.
		// IsContextError returns false for such errors, so the old code fell through
		// to transport wrapping even when the stream context was already cancelled.
		streamErr := errors.New("stream transport: connection closed")
		mockStream.EXPECT().CloseSend().Return(streamErr)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		cancelledCtx, cancel := context.WithCancel(t.Context())
		cancel() // cancel the stream context before calling CloseSend

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  cancelledCtx,
			wrapping:   true,
			traceID:    "test-trace-id",
		}

		err := s.CloseSend()
		require.Error(t, err)
		// When the stream context is done, the error must NOT be wrapped as a
		// transport error regardless of what gRPC returned.
		require.True(t, xerrors.IsTransportError(err))
		require.ErrorIs(t, err, streamErr)
	})

	t.Run("TransportErrorWithWrapping", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		mockStream.EXPECT().CloseSend().Return(fmt.Errorf("transport error"))

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135", endpoint.WithID(123))
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   true,
			traceID:    "test-trace-id",
		}

		err := s.CloseSend()
		require.Error(t, err)
		require.True(t, xerrors.IsTransportError(err))
	})

	t.Run("ErrorWithoutWrapping", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		expectedErr := fmt.Errorf("raw error")
		mockStream.EXPECT().CloseSend().Return(expectedErr)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   false,
		}

		err := s.CloseSend()
		require.Error(t, err)
		require.Equal(t, expectedErr, err)
	})
}

func TestGrpcClientStream_SendMsg(t *testing.T) {
	t.Run("SuccessfulSend", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryRequest{}
		mockStream.EXPECT().SendMsg(msg).Return(nil)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   true,
			sentMark:   &modificationMark{},
		}

		err := s.SendMsg(msg)
		require.NoError(t, err)
	})

	t.Run("ContextError", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryRequest{}
		mockStream.EXPECT().SendMsg(msg).Return(context.DeadlineExceeded)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   true,
			sentMark:   &modificationMark{},
		}

		err := s.SendMsg(msg)
		require.Error(t, err)
		require.True(t, xerrors.IsContextError(err))
	})

	t.Run("StreamContextDoneReturnsNonTransportError", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryRequest{}
		// Use a non-gRPC-status error that gRPC may return on stream termination.
		// IsContextError returns false for such errors, so the old code fell through
		// to transport wrapping even when the stream context was already cancelled.
		streamErr := errors.New("stream transport: connection closed")
		mockStream.EXPECT().SendMsg(msg).Return(streamErr)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		cancelledCtx, cancel := context.WithCancel(t.Context())
		cancel() // cancel the stream context before calling SendMsg

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  cancelledCtx,
			wrapping:   true,
			traceID:    "test-trace-id",
			sentMark:   &modificationMark{},
		}

		err := s.SendMsg(msg)
		require.Error(t, err)
		// When the stream context is done, the error must NOT be wrapped as a
		// transport error regardless of what gRPC returned.
		require.True(t, xerrors.IsTransportError(err))
		require.ErrorIs(t, err, streamErr)
	})

	t.Run("TransportErrorRetryable", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryRequest{}
		mockStream.EXPECT().SendMsg(msg).Return(grpcStatus.Error(grpcCodes.Unavailable, "unavailable"))

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   true,
			traceID:    "test-trace-id",
			sentMark:   &modificationMark{},
		}

		err := s.SendMsg(msg)
		require.Error(t, err)
		require.True(t, xerrors.IsTransportError(err))
		require.True(t, xerrors.IsRetryableError(err))
	})

	t.Run("TransportErrorNonRetryable", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryRequest{}
		mockStream.EXPECT().SendMsg(msg).Return(grpcStatus.Error(grpcCodes.Unavailable, "unavailable"))

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135", endpoint.WithID(123))
		parentConn := newConn(e, config)

		mark := &modificationMark{}
		mark.markDirty()

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   true,
			traceID:    "test-trace-id",
			sentMark:   mark,
		}

		err := s.SendMsg(msg)
		require.Error(t, err)
		require.True(t, xerrors.IsTransportError(err))
		require.False(t, xerrors.IsRetryableError(err))
	})

	t.Run("ErrorWithoutWrapping", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryRequest{}
		expectedErr := fmt.Errorf("raw error")
		mockStream.EXPECT().SendMsg(msg).Return(expectedErr)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   false,
			sentMark:   &modificationMark{},
		}

		err := s.SendMsg(msg)
		require.Error(t, err)
		require.Equal(t, expectedErr, err)
	})
}

func TestGrpcClientStream_RecvMsg(t *testing.T) {
	t.Run("SuccessfulReceive", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryResponsePart{}
		mockStream.EXPECT().RecvMsg(msg).DoAndReturn(func(m any) error {
			resp := m.(*Ydb_Query.ExecuteQueryResponsePart)
			resp.Status = Ydb.StatusIds_SUCCESS

			return nil
		})

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   true,
			sentMark:   &modificationMark{},
		}

		err := s.RecvMsg(msg)
		require.NoError(t, err)
	})

	t.Run("EOFError", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryResponsePart{}
		mockStream.EXPECT().RecvMsg(msg).Return(io.EOF)
		mockStream.EXPECT().Trailer().Return(metadata.MD{})

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   true,
			sentMark:   &modificationMark{},
		}

		err := s.RecvMsg(msg)
		require.Equal(t, io.EOF, err)
	})

	t.Run("ContextError", func(t *testing.T) {
		t.Run("context.Canceled from RecvMsg", func(t *testing.T) {
			ctrl := gomock.NewController(t)
			mockStream := mock.NewMockClientStream(ctrl)

			msg := &Ydb_Query.ExecuteQueryResponsePart{}
			mockStream.EXPECT().RecvMsg(msg).Return(context.Canceled)
			mockStream.EXPECT().Trailer().Return(metadata.MD{})

			config := &mockConfig{
				dialTimeout:   5 * time.Second,
				connectionTTL: 0,
			}
			e := endpoint.New("test-endpoint:2135")
			parentConn := newConn(e, config)

			s := &grpcClientStream{
				parentConn: parentConn,
				stream:     mockStream,
				streamCtx:  t.Context(),
				wrapping:   true,
				sentMark:   &modificationMark{},
			}

			err := s.RecvMsg(msg)
			require.Error(t, err)
			require.True(t, xerrors.IsContextError(err))
		})
		t.Run("stream context canceled", func(t *testing.T) {
			ctrl := gomock.NewController(t)
			mockStream := mock.NewMockClientStream(ctrl)

			msg := &Ydb_Query.ExecuteQueryResponsePart{}
			mockStream.EXPECT().RecvMsg(msg).Return(grpcStatus.Error(grpcCodes.Canceled, ""))
			mockStream.EXPECT().Trailer().Return(metadata.MD{})

			config := &mockConfig{
				dialTimeout:   5 * time.Second,
				connectionTTL: 0,
			}
			e := endpoint.New("test-endpoint:2135")
			parentConn := newConn(e, config)

			ctx, cancel := context.WithCancel(t.Context())
			cancel()

			s := &grpcClientStream{
				parentConn: parentConn,
				stream:     mockStream,
				streamCtx:  ctx,
				wrapping:   true,
				sentMark:   &modificationMark{},
			}

			err := s.RecvMsg(msg)
			require.Error(t, err)
			require.True(t, xerrors.IsContextError(err))
			require.True(t, xerrors.IsTransportError(err))
		})
	})

	t.Run("StreamContextDoneReturnsNonTransportError", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryResponsePart{}
		// Use a non-gRPC-status error that gRPC may return on stream termination.
		// IsContextError returns false for such errors, so the old code fell through
		// to transport wrapping even when the stream context was already cancelled.
		streamErr := errors.New("stream transport: connection closed")
		mockStream.EXPECT().RecvMsg(msg).Return(streamErr)
		mockStream.EXPECT().Trailer().Return(metadata.MD{})

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		cancelledCtx, cancel := context.WithCancel(t.Context())
		cancel() // cancel the stream context before calling RecvMsg

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  cancelledCtx,
			wrapping:   true,
			traceID:    "test-trace-id",
			sentMark:   &modificationMark{},
		}

		err := s.RecvMsg(msg)
		require.Error(t, err)
		// When the stream context is done, the error must NOT be wrapped as a
		// transport error regardless of what gRPC returned.
		require.True(t, xerrors.IsTransportError(err))
		require.True(t, xerrors.IsContextError(err))
		require.ErrorIs(t, err, streamErr)
	})

	t.Run("TransportErrorRetryable", func(t *testing.T) {
		t.Run("Unavailable", func(t *testing.T) {
			ctrl := gomock.NewController(t)
			mockStream := mock.NewMockClientStream(ctrl)

			msg := &Ydb_Query.ExecuteQueryResponsePart{}
			mockStream.EXPECT().RecvMsg(msg).Return(grpcStatus.Error(grpcCodes.Unavailable, "unavailable"))
			mockStream.EXPECT().Trailer().Return(metadata.MD{})

			config := &mockConfig{
				dialTimeout:   5 * time.Second,
				connectionTTL: 0,
			}
			e := endpoint.New("test-endpoint:2135")
			parentConn := newConn(e, config)

			s := &grpcClientStream{
				parentConn: parentConn,
				stream:     mockStream,
				streamCtx:  t.Context(),
				wrapping:   true,
				traceID:    "test-trace-id",
				sentMark:   &modificationMark{},
			}

			err := s.RecvMsg(msg)
			require.Error(t, err)
			require.True(t, xerrors.IsTransportError(err))
			require.True(t, xerrors.IsRetryableError(err))
		})
		t.Run("Cancelled", func(t *testing.T) {
			t.Run("sentMark.canRetry()==false", func(t *testing.T) {
				ctrl := gomock.NewController(t)
				mockStream := mock.NewMockClientStream(ctrl)

				msg := &Ydb_Query.ExecuteQueryResponsePart{}
				mockStream.EXPECT().RecvMsg(msg).Return(grpcStatus.Error(grpcCodes.Canceled, "Cancelled on the server side"))
				mockStream.EXPECT().Trailer().Return(metadata.MD{})

				ctx, cancel := context.WithCancel(t.Context())
				cancel()

				// Cleaner existing pattern:
				mark := &modificationMark{}
				mark.markDirty()

				s := &grpcClientStream{
					parentConn: &conn{
						config:    &mockConfig{},
						endpoint:  endpoint.New("test-endpoint:2135"),
						lastUsage: nopLastUsage{},
					},
					stream:    mockStream,
					streamCtx: ctx,
					wrapping:  true,
					sentMark:  mark,
				}

				err := s.RecvMsg(msg)
				require.Error(t, err)

				check := retry.Check(err)
				require.EqualValues(t, grpcCodes.Canceled, check.StatusCode())
				require.EqualValues(t, backoff.TypeFast, check.BackoffType())
				require.True(t, check.MustRetry(true))
				require.False(t, check.MustRetry(false))
			})
			t.Run("sentMark.canRetry()==true", func(t *testing.T) {
				ctrl := gomock.NewController(t)
				mockStream := mock.NewMockClientStream(ctrl)

				msg := &Ydb_Query.ExecuteQueryResponsePart{}
				mockStream.EXPECT().RecvMsg(msg).Return(grpcStatus.Error(grpcCodes.Canceled, context.Canceled.Error()))
				mockStream.EXPECT().Trailer().Return(metadata.MD{})

				ctx, cancel := context.WithCancel(t.Context())
				cancel()

				s := &grpcClientStream{
					parentConn: &conn{
						config:    &mockConfig{},
						endpoint:  endpoint.New("test-endpoint:2135"),
						lastUsage: nopLastUsage{},
					},
					stream:    mockStream,
					streamCtx: ctx,
					wrapping:  true,
					sentMark:  &modificationMark{},
				}

				err := s.RecvMsg(msg)
				require.Error(t, err)

				check := retry.Check(err)
				require.EqualValues(t, grpcCodes.Canceled, check.StatusCode())
				require.EqualValues(t, backoff.TypeFast, check.BackoffType())
				require.True(t, check.MustRetry(true))
				require.True(t, check.MustRetry(false))
			})
		})
	})

	t.Run("TransportErrorNonRetryable", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryResponsePart{}
		mockStream.EXPECT().RecvMsg(msg).Return(grpcStatus.Error(grpcCodes.Unavailable, "unavailable"))
		mockStream.EXPECT().Trailer().Return(metadata.MD{})

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135", endpoint.WithID(123))
		parentConn := newConn(e, config)

		mark := &modificationMark{}
		mark.markDirty()

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   true,
			traceID:    "test-trace-id",
			sentMark:   mark,
		}

		err := s.RecvMsg(msg)
		require.Error(t, err)
		require.True(t, xerrors.IsTransportError(err))
		require.False(t, xerrors.IsRetryableError(err))
	})

	t.Run("ErrorWithoutWrapping", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryResponsePart{}
		expectedErr := fmt.Errorf("raw error")
		mockStream.EXPECT().RecvMsg(msg).Return(expectedErr)
		mockStream.EXPECT().Trailer().Return(metadata.MD{})

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   false,
			sentMark:   &modificationMark{},
		}

		err := s.RecvMsg(msg)
		require.Error(t, err)
		require.Equal(t, expectedErr, err)
	})

	t.Run("OperationError", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryResponsePart{}
		mockStream.EXPECT().RecvMsg(msg).DoAndReturn(func(m any) error {
			resp := m.(*Ydb_Query.ExecuteQueryResponsePart)
			resp.Status = Ydb.StatusIds_UNAVAILABLE

			return nil
		})
		mockStream.EXPECT().Trailer().Return(metadata.MD{})

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135", endpoint.WithID(123))
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   true,
			sentMark:   &modificationMark{},
		}

		err := s.RecvMsg(msg)
		require.Error(t, err)
		require.True(t, xerrors.IsOperationError(err, Ydb.StatusIds_UNAVAILABLE))
	})

	t.Run("OperationErrorWithoutWrapping", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		msg := &Ydb_Query.ExecuteQueryResponsePart{}
		mockStream.EXPECT().RecvMsg(msg).DoAndReturn(func(m any) error {
			resp := m.(*Ydb_Query.ExecuteQueryResponsePart)
			resp.Status = Ydb.StatusIds_UNAVAILABLE

			return nil
		})

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		s := &grpcClientStream{
			parentConn: parentConn,
			stream:     mockStream,
			streamCtx:  t.Context(),
			wrapping:   false,
			sentMark:   &modificationMark{},
		}

		err := s.RecvMsg(msg)
		require.NoError(t, err)
	})
}

func TestGrpcClientStream_Finish(t *testing.T) {
	t.Run("CallsCancelOnFinish", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
			driverTrace:   &trace.Driver{},
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		ctx, cancel := context.WithCancel(t.Context())
		called := false
		wrappedCancel := func() {
			called = true
			cancel()
		}

		s := &grpcClientStream{
			parentConn:   parentConn,
			stream:       mockStream,
			streamCtx:    ctx,
			streamCancel: wrappedCancel,
		}

		s.finish(nil)
		require.True(t, called)
	})

	t.Run("FinishWithError", func(t *testing.T) {
		ctrl := gomock.NewController(t)
		mockStream := mock.NewMockClientStream(ctrl)

		config := &mockConfig{
			dialTimeout:   5 * time.Second,
			connectionTTL: 0,
			driverTrace:   &trace.Driver{},
		}
		e := endpoint.New("test-endpoint:2135")
		parentConn := newConn(e, config)

		ctx, cancel := context.WithCancel(t.Context())

		s := &grpcClientStream{
			parentConn:   parentConn,
			stream:       mockStream,
			streamCtx:    ctx,
			streamCancel: cancel,
		}

		testErr := fmt.Errorf("test error")
		s.finish(testErr)
		// Should not panic
	})
}
