github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/topic/topicwriterinternal/writer_grpc_mock_test.go (about)

     1  package topicwriterinternal_test
     2  
     3  import (
     4  	"errors"
     5  	"fmt"
     6  	"strings"
     7  	"testing"
     8  
     9  	"github.com/rekby/fixenv"
    10  	"github.com/rekby/fixenv/sf"
    11  	"github.com/stretchr/testify/require"
    12  	"github.com/ydb-platform/ydb-go-genproto/Ydb_Topic_V1"
    13  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
    14  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Issue"
    15  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Topic"
    16  
    17  	"github.com/ydb-platform/ydb-go-sdk/v3"
    18  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xtest"
    19  	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topicoptions"
    20  	"github.com/ydb-platform/ydb-go-sdk/v3/topic/topicwriter"
    21  )
    22  
    23  func TestRegressionOperationUnavailableIssue1007(t *testing.T) {
    24  	xtest.TestManyTimes(t, func(t testing.TB) {
    25  		e := fixenv.New(t)
    26  
    27  		mock := newTopicWriterOperationUnavailable()
    28  		connString := xtest.GrpcMockTopicConnString(e, mock)
    29  
    30  		db, err := ydb.Open(sf.Context(e), connString)
    31  		require.NoError(t, err)
    32  
    33  		writer, err := db.Topic().StartWriter("test", topicoptions.WithWriterWaitServerAck(true))
    34  		require.NoError(t, err)
    35  
    36  		err = writer.Write(sf.Context(e), topicwriter.Message{
    37  			Data: strings.NewReader("asd"),
    38  		})
    39  		require.NoError(t, err)
    40  		require.True(t, mock.UnavailableResponsed)
    41  	})
    42  }
    43  
    44  type topicWriterOperationUnavailable struct {
    45  	Ydb_Topic_V1.UnimplementedTopicServiceServer
    46  
    47  	UnavailableResponsed bool
    48  }
    49  
    50  func newTopicWriterOperationUnavailable() *topicWriterOperationUnavailable {
    51  	return &topicWriterOperationUnavailable{}
    52  }
    53  
    54  func (t *topicWriterOperationUnavailable) StreamWrite(server Ydb_Topic_V1.TopicService_StreamWriteServer) error {
    55  	initMsg, err := server.Recv()
    56  	if err != nil {
    57  		return fmt.Errorf("failed read init message: %w", err)
    58  	}
    59  
    60  	if initMsg.GetInitRequest() == nil {
    61  		return errors.New("first message must be init message")
    62  	}
    63  
    64  	err = server.Send(&Ydb_Topic.StreamWriteMessage_FromServer{
    65  		Status: Ydb.StatusIds_SUCCESS,
    66  		ServerMessage: &Ydb_Topic.StreamWriteMessage_FromServer_InitResponse{
    67  			InitResponse: &Ydb_Topic.StreamWriteMessage_InitResponse{
    68  				LastSeqNo:       0,
    69  				SessionId:       "test",
    70  				PartitionId:     0,
    71  				SupportedCodecs: nil,
    72  			},
    73  		},
    74  	})
    75  	if err != nil {
    76  		return fmt.Errorf("failed to send init response: %w", err)
    77  	}
    78  
    79  	if !t.UnavailableResponsed {
    80  		t.UnavailableResponsed = true
    81  
    82  		err = server.Send(&Ydb_Topic.StreamWriteMessage_FromServer{
    83  			Status: Ydb.StatusIds_UNAVAILABLE,
    84  			Issues: []*Ydb_Issue.IssueMessage{
    85  				{
    86  					Message: "Test status unavailable",
    87  				},
    88  			},
    89  		})
    90  
    91  		if err != nil {
    92  			return fmt.Errorf("failed to send error response: %w", err)
    93  		}
    94  
    95  		return nil
    96  	}
    97  
    98  	// wait message block
    99  	messagesMsg, err := server.Recv()
   100  	if err != nil {
   101  		return errors.New("failed to read messages block")
   102  	}
   103  
   104  	if len(messagesMsg.GetClientMessage().(*Ydb_Topic.StreamWriteMessage_FromClient_WriteRequest).
   105  		WriteRequest.GetMessages()) == 0 {
   106  		return errors.New("received zero messages block")
   107  	}
   108  
   109  	err = server.Send(&Ydb_Topic.StreamWriteMessage_FromServer{
   110  		Status: Ydb.StatusIds_SUCCESS,
   111  		ServerMessage: &Ydb_Topic.StreamWriteMessage_FromServer_WriteResponse{
   112  			WriteResponse: &Ydb_Topic.StreamWriteMessage_WriteResponse{
   113  				Acks: []*Ydb_Topic.StreamWriteMessage_WriteResponse_WriteAck{
   114  					{
   115  						SeqNo: 1,
   116  						MessageWriteStatus: &Ydb_Topic.StreamWriteMessage_WriteResponse_WriteAck_Written_{
   117  							Written: &Ydb_Topic.StreamWriteMessage_WriteResponse_WriteAck_Written{
   118  								Offset: 1,
   119  							},
   120  						},
   121  					},
   122  				},
   123  				PartitionId:     0,
   124  				WriteStatistics: &Ydb_Topic.StreamWriteMessage_WriteResponse_WriteStatistics{},
   125  			},
   126  		},
   127  	})
   128  
   129  	if err != nil {
   130  		return fmt.Errorf("failed to sent write ack: %w", err)
   131  	}
   132  
   133  	return nil
   134  }