github.com/ronaksoft/rony@v0.16.26-0.20230807065236-1743dbfe6959/edgetest/context_rpc.go (about)

     1  package edgetest
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	"time"
     7  
     8  	"github.com/ronaksoft/rony"
     9  	dummyGateway "github.com/ronaksoft/rony/internal/gateway/dummy"
    10  	"github.com/ronaksoft/rony/tools"
    11  	"google.golang.org/protobuf/proto"
    12  )
    13  
    14  /*
    15     Creation Time: 2020 - Dec - 09
    16     Created by:  (ehsan)
    17     Maintainers:
    18        1.  Ehsan N. Moosa (E2)
    19     Auditor: Ehsan N. Moosa (E2)
    20     Copyright Ronak Software Group 2020
    21  */
    22  
    23  type rpcCtx struct {
    24  	mtx        sync.Mutex
    25  	id         uint64
    26  	reqC       uint64
    27  	reqID      uint64
    28  	req        []byte
    29  	expect     map[uint64]CheckFunc
    30  	gw         *dummyGateway.Gateway
    31  	err        error
    32  	errH       func(constructor uint64, e *rony.Error)
    33  	doneCh     chan struct{}
    34  	kvs        []*rony.KeyValue
    35  	persistent bool
    36  }
    37  
    38  func newRPCContext(gw *dummyGateway.Gateway) *rpcCtx {
    39  	c := &rpcCtx{
    40  		id:     atomic.AddUint64(&connID, 1),
    41  		expect: make(map[uint64]CheckFunc),
    42  		gw:     gw,
    43  		doneCh: make(chan struct{}, 1),
    44  	}
    45  
    46  	return c
    47  }
    48  
    49  // Persistent makes the rpcCtx simulate persistent connection e.g. websocket
    50  func (c *rpcCtx) Persistent() *rpcCtx {
    51  	c.persistent = true
    52  
    53  	return c
    54  }
    55  
    56  // Request set the request you wish to send to the server
    57  func (c *rpcCtx) Request(constructor uint64, p rony.IMessage, kvs ...*rony.KeyValue) *rpcCtx {
    58  	c.reqID = tools.RandomUint64(0)
    59  	c.reqC = constructor
    60  	data, _ := proto.Marshal(p)
    61  	e := &rony.MessageEnvelope{
    62  		Constructor: constructor,
    63  		RequestID:   c.reqID,
    64  		Message:     data,
    65  		Auth:        nil,
    66  		Header:      kvs,
    67  	}
    68  
    69  	c.req, c.err = proto.Marshal(e)
    70  
    71  	return c
    72  }
    73  
    74  // Expect let you set what you expect to receive. If cf is set, then you can do more checks
    75  // on the response and return error if the response was not fully acceptable
    76  func (c *rpcCtx) Expect(constructor uint64, cf CheckFunc) *rpcCtx {
    77  	c.expect[constructor] = cf
    78  
    79  	return c
    80  }
    81  
    82  func (c *rpcCtx) ExpectConstructor(constructor uint64) *rpcCtx {
    83  	return c.Expect(constructor, nil)
    84  }
    85  
    86  func (c *rpcCtx) check(e *rony.MessageEnvelope) {
    87  	c.mtx.Lock()
    88  	f, ok := c.expect[e.Constructor]
    89  	c.mtx.Unlock()
    90  	if !ok && e.Constructor == rony.C_Error {
    91  		err := &rony.Error{}
    92  		c.err = err.Unmarshal(e.Message)
    93  		if c.errH != nil {
    94  			c.errH(c.reqC, err)
    95  		}
    96  
    97  		return
    98  	}
    99  	if f != nil {
   100  		c.err = f(e.Message, e.Header...)
   101  	}
   102  	c.mtx.Lock()
   103  	delete(c.expect, e.Constructor)
   104  	c.mtx.Unlock()
   105  }
   106  
   107  func (c *rpcCtx) expectCount() int {
   108  	c.mtx.Lock()
   109  	n := len(c.expect)
   110  	c.mtx.Unlock()
   111  
   112  	return n
   113  }
   114  
   115  func (c *rpcCtx) ErrorHandler(f func(constructor uint64, e *rony.Error)) *rpcCtx {
   116  	c.errH = f
   117  
   118  	return c
   119  }
   120  
   121  func (c *rpcCtx) SetRunParameters(kvs ...*rony.KeyValue) *rpcCtx {
   122  	c.kvs = kvs
   123  
   124  	return c
   125  }
   126  
   127  func (c *rpcCtx) RunShort(kvs ...*rony.KeyValue) error {
   128  	return c.Run(time.Second*10, kvs...)
   129  }
   130  
   131  func (c *rpcCtx) RunLong(kvs ...*rony.KeyValue) error {
   132  	return c.Run(time.Minute, kvs...)
   133  }
   134  
   135  func (c *rpcCtx) Run(timeout time.Duration, kvs ...*rony.KeyValue) error {
   136  	// We return error early if we have encountered error before Run
   137  	if c.err != nil {
   138  		return c.err
   139  	}
   140  
   141  	c.SetRunParameters(kvs...)
   142  
   143  	// Open Connection
   144  	c.gw.OpenConn(c.id, c.persistent, c.receiver, c.kvs...)
   145  
   146  	// Send the Request
   147  	err := c.gw.RPC(c.id, 0, c.req)
   148  	if err != nil {
   149  		return err
   150  	}
   151  
   152  	// Wait for Response(s)
   153  Loop:
   154  	for {
   155  		select {
   156  		case <-c.doneCh:
   157  			// Check if all the expectations have been passed
   158  			if c.expectCount() == 0 {
   159  				break Loop
   160  			}
   161  		case <-time.After(timeout):
   162  			break Loop
   163  		}
   164  	}
   165  
   166  	if c.expectCount() > 0 {
   167  		c.err = ErrExpectationFailed
   168  	}
   169  
   170  	return c.err
   171  }
   172  
   173  func (c *rpcCtx) receiver(connID uint64, streamID int64, data []byte) {
   174  	defer func() {
   175  		c.doneCh <- struct{}{}
   176  	}()
   177  	e := &rony.MessageEnvelope{}
   178  	c.err = e.Unmarshal(data)
   179  	if c.err != nil {
   180  		return
   181  	}
   182  	switch e.Constructor {
   183  	case rony.C_MessageContainer:
   184  		mc := &rony.MessageContainer{}
   185  		c.err = mc.Unmarshal(e.Message)
   186  		if c.err != nil {
   187  			return
   188  		}
   189  		for _, e := range mc.Envelopes {
   190  			c.check(e)
   191  		}
   192  	default:
   193  		c.check(e)
   194  	}
   195  }