github.com/tickoalcantara12/micro/v3@v3.0.0-20221007104245-9d75b9bcbab9/service/server/mucp/rpc_stream_test.go (about)

     1  // Copyright 2020 Asim Aslam
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  // Original source: github.com/micro/go-micro/v3/server/mucp/rpc_stream_test.go
    16  
    17  package mucp
    18  
    19  import (
    20  	"bytes"
    21  	"fmt"
    22  	"io"
    23  	"math/rand"
    24  	"sync"
    25  	"testing"
    26  	"time"
    27  
    28  	"github.com/golang/protobuf/proto"
    29  	"github.com/tickoalcantara12/micro/v3/util/codec/json"
    30  	protoCodec "github.com/tickoalcantara12/micro/v3/util/codec/proto"
    31  )
    32  
    33  // protoStruct implements proto.Message
    34  type protoStruct struct {
    35  	Payload string `protobuf:"bytes,1,opt,name=service,proto3" json:"service,omitempty"`
    36  }
    37  
    38  func (m *protoStruct) Reset()         { *m = protoStruct{} }
    39  func (m *protoStruct) String() string { return proto.CompactTextString(m) }
    40  func (*protoStruct) ProtoMessage()    {}
    41  
    42  // safeBuffer throws away everything and wont Read data back
    43  type safeBuffer struct {
    44  	sync.RWMutex
    45  	buf []byte
    46  	off int
    47  }
    48  
    49  func (b *safeBuffer) Write(p []byte) (n int, err error) {
    50  	if len(p) == 0 {
    51  		return 0, nil
    52  	}
    53  	// Cannot retain p, so we must copy it:
    54  	p2 := make([]byte, len(p))
    55  	copy(p2, p)
    56  	b.Lock()
    57  	b.buf = append(b.buf, p2...)
    58  	b.Unlock()
    59  	return len(p2), nil
    60  }
    61  
    62  func (b *safeBuffer) Read(p []byte) (n int, err error) {
    63  	if len(p) == 0 {
    64  		return 0, nil
    65  	}
    66  	b.RLock()
    67  	n = copy(p, b.buf[b.off:])
    68  	b.RUnlock()
    69  	if n == 0 {
    70  		return 0, io.EOF
    71  	}
    72  	b.off += n
    73  	return n, nil
    74  }
    75  
    76  func (b *safeBuffer) Close() error {
    77  	return nil
    78  }
    79  
    80  func TestRPCStream_Sequence(t *testing.T) {
    81  	buffer := new(bytes.Buffer)
    82  	rwc := readWriteCloser{
    83  		rbuf: buffer,
    84  		wbuf: buffer,
    85  	}
    86  	codec := json.NewCodec(&rwc)
    87  	streamServer := rpcStream{
    88  		codec: codec,
    89  		request: &rpcRequest{
    90  			codec: codec,
    91  		},
    92  	}
    93  
    94  	// Check if sequence is correct
    95  	for i := 0; i < 1000; i++ {
    96  		if err := streamServer.Send(fmt.Sprintf(`{"test":"value %d"}`, i)); err != nil {
    97  			t.Errorf("Unexpected Send error: %s", err)
    98  		}
    99  	}
   100  
   101  	for i := 0; i < 1000; i++ {
   102  		var msg string
   103  		if err := streamServer.Recv(&msg); err != nil {
   104  			t.Errorf("Unexpected Recv error: %s", err)
   105  		}
   106  		if msg != fmt.Sprintf(`{"test":"value %d"}`, i) {
   107  			t.Errorf("Unexpected msg: %s", msg)
   108  		}
   109  	}
   110  }
   111  
   112  func TestRPCStream_Concurrency(t *testing.T) {
   113  	buffer := new(safeBuffer)
   114  	codec := protoCodec.NewCodec(buffer)
   115  	streamServer := rpcStream{
   116  		codec: codec,
   117  		request: &rpcRequest{
   118  			codec: codec,
   119  		},
   120  	}
   121  
   122  	var wg sync.WaitGroup
   123  	// Check if race conditions happen
   124  	for i := 0; i < 10; i++ {
   125  		wg.Add(2)
   126  
   127  		go func() {
   128  			for i := 0; i < 50; i++ {
   129  				msg := protoStruct{Payload: "test"}
   130  				<-time.After(time.Duration(rand.Intn(50)) * time.Millisecond)
   131  				if err := streamServer.Send(msg); err != nil {
   132  					t.Errorf("Unexpected Send error: %s", err)
   133  				}
   134  			}
   135  			wg.Done()
   136  		}()
   137  
   138  		go func() {
   139  			for i := 0; i < 50; i++ {
   140  				<-time.After(time.Duration(rand.Intn(50)) * time.Millisecond)
   141  				if err := streamServer.Recv(&protoStruct{}); err != nil {
   142  					t.Errorf("Unexpected Recv error: %s", err)
   143  				}
   144  			}
   145  			wg.Done()
   146  		}()
   147  	}
   148  	wg.Wait()
   149  }