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

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