github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/test/mock/test_server.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     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  //     http://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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package mock
    15  
    16  import (
    17  	"context"
    18  	"io"
    19  	"sync/atomic"
    20  
    21  	pb "github.com/pingcap/tiflow/engine/enginepb"
    22  	"github.com/pingcap/tiflow/pkg/errors"
    23  	"google.golang.org/grpc"
    24  	"google.golang.org/grpc/metadata"
    25  )
    26  
    27  type testServer struct {
    28  	*baseServer
    29  	pb.TestServiceServer
    30  }
    31  
    32  type testServerConn struct {
    33  	server *testServer
    34  	stream *testStream
    35  }
    36  
    37  func (s *testServer) dial() (Conn, error) {
    38  	return &testServerConn{s, nil}, nil
    39  }
    40  
    41  type testClient struct {
    42  	conn Conn
    43  }
    44  
    45  func (t *testClient) FeedBinlog(ctx context.Context, in *pb.TestBinlogRequest, opts ...grpc.CallOption) (pb.TestService_FeedBinlogClient, error) {
    46  	resp, err := t.conn.sendRequest(ctx, in)
    47  	return resp.(pb.TestService_FeedBinlogClient), err
    48  }
    49  
    50  func (s *testServerConn) Close() error {
    51  	return nil
    52  }
    53  
    54  type testStream struct {
    55  	ctx    context.Context
    56  	data   chan *pb.Record
    57  	err    error
    58  	closed int32
    59  }
    60  
    61  func (s *testStream) close() {
    62  	if atomic.CompareAndSwapInt32(&s.closed, 0, 1) {
    63  		s.err = io.EOF
    64  		close(s.data)
    65  	}
    66  }
    67  
    68  func (s *testStream) Send(r *pb.Record) error {
    69  	if atomic.LoadInt32(&s.closed) == 1 {
    70  		return errors.New("stream has been closed")
    71  	}
    72  	s.data <- r
    73  	return nil
    74  }
    75  
    76  func (s *testStream) Recv() (*pb.Record, error) {
    77  	select {
    78  	case r := <-s.data:
    79  		if r == nil {
    80  			return nil, s.err
    81  		}
    82  		return r, nil
    83  	case <-s.ctx.Done():
    84  		return nil, errors.New("canceled")
    85  	}
    86  }
    87  
    88  func (s *testStream) SetHeader(metadata.MD) error {
    89  	return errors.New("unimplemented")
    90  }
    91  
    92  func (s *testStream) SendHeader(metadata.MD) error {
    93  	return errors.New("unimplemented")
    94  }
    95  
    96  func (s *testStream) SetTrailer(metadata.MD) {}
    97  
    98  func (s *testStream) Context() context.Context { return nil }
    99  
   100  func (s *testStream) SendMsg(interface{}) error {
   101  	return errors.New("unimplemented")
   102  }
   103  
   104  func (s *testStream) RecvMsg(interface{}) error {
   105  	return errors.New("unimplemented")
   106  }
   107  
   108  func (s *testStream) Header() (metadata.MD, error) {
   109  	return nil, errors.New("unimplemented")
   110  }
   111  
   112  func (s *testStream) Trailer() metadata.MD { return nil }
   113  
   114  func (s *testStream) CloseSend() error {
   115  	return errors.New("unimplemented")
   116  }
   117  
   118  func (s *testServerConn) sendRequest(ctx context.Context, req interface{}) (interface{}, error) {
   119  	switch x := req.(type) {
   120  	case *pb.TestBinlogRequest:
   121  		stream := &testStream{
   122  			data: make(chan *pb.Record, 1024),
   123  			ctx:  ctx,
   124  		}
   125  		if s.stream != nil {
   126  			return nil, errors.New("internal error")
   127  		}
   128  		s.stream = stream
   129  		go func() {
   130  			stream.err = s.server.FeedBinlog(x, stream)
   131  			stream.close()
   132  		}()
   133  		return stream, nil
   134  	}
   135  	return nil, errors.New("unknown request")
   136  }
   137  
   138  // NewTestServer creates a mock gRPC server
   139  func NewTestServer(addr string, server pb.TestServiceServer) (GrpcServer, error) {
   140  	container.mu.Lock()
   141  	defer container.mu.Unlock()
   142  	_, ok := container.servers[addr]
   143  	if ok {
   144  		return nil, errors.New("addr " + addr + " has been listened")
   145  	}
   146  	newServer := &testServer{&baseServer{addr}, server}
   147  	container.servers[addr] = newServer
   148  	return newServer, nil
   149  }
   150  
   151  // NewTestClient creates a mock test service client
   152  func NewTestClient(conn Conn) pb.TestServiceClient {
   153  	return &testClient{conn}
   154  }