woop/p2p/stream/common/requestmanager/interface_test.go

package requestmanager

import (
	"errors"
	"fmt"
	"strconv"
	"sync"

	"github.com/ethereum/go-ethereum/event"
	"github.com/ethereum/go-ethereum/rlp"
	"github.com/harmony-one/harmony/p2p/stream/common/streammanager"
	sttypes "github.com/harmony-one/harmony/p2p/stream/types"
)

var testProtoID = sttypes.ProtoID("harmony/sync/unitest/0/1.0.0")

type testStreamManager struct {
	streams map[sttypes.StreamID]sttypes.Stream

	newStreamFeed event.Feed
	rmStreamFeed  event.Feed

	lock sync.Mutex
}

func newTestStreamManager() *testStreamManager {
	return &testStreamManager{
		streams: make(map[sttypes.StreamID]sttypes.Stream),
	}
}

func (sm *testStreamManager) addNewStream(st sttypes.Stream) {
	sm.lock.Lock()
	sm.streams[st.ID()] = st
	sm.lock.Unlock()

	sm.newStreamFeed.Send(streammanager.EvtStreamAdded{Stream: st})
}

func (sm *testStreamManager) rmStream(stid sttypes.StreamID) {
	sm.lock.Lock()
	delete(sm.streams, stid)
	sm.lock.Unlock()

	sm.rmStreamFeed.Send(streammanager.EvtStreamRemoved{ID: stid})
}

func (sm *testStreamManager) SubscribeAddStreamEvent(ch chan<- streammanager.EvtStreamAdded) event.Subscription {
	return sm.newStreamFeed.Subscribe(ch)
}

func (sm *testStreamManager) SubscribeRemoveStreamEvent(ch chan<- streammanager.EvtStreamRemoved) event.Subscription {
	return sm.rmStreamFeed.Subscribe(ch)
}

func (sm *testStreamManager) GetStreams() []sttypes.Stream {
	sm.lock.Lock()
	defer sm.lock.Unlock()

	sts := make([]sttypes.Stream, 0, len(sm.streams))

	for _, st := range sm.streams {
		sts = append(sts, st)
	}
	return sts
}

func (sm *testStreamManager) GetStreamByID(id sttypes.StreamID) (sttypes.Stream, bool) {
	sm.lock.Lock()
	defer sm.lock.Unlock()

	st, exist := sm.streams[id]
	return st, exist
}

type testStream struct {
	id      sttypes.StreamID
	rm      *requestManager
	deliver func(*testRequest) // use goroutine inside this function
}

func (st *testStream) ID() sttypes.StreamID {
	return st.id
}

func (st *testStream) ProtoID() sttypes.ProtoID {
	return testProtoID
}

func (st *testStream) WriteBytes(b []byte) error {
	req, err := decodeTestRequest(b)
	if err != nil {
		return err
	}
	if st.rm != nil && st.deliver != nil {
		st.deliver(req)
	}
	return nil
}

func (st *testStream) ReadBytes() ([]byte, error) {
	return nil, nil
}

func (st *testStream) ProtoSpec() (sttypes.ProtoSpec, error) {
	return sttypes.ProtoIDToProtoSpec(testProtoID)
}

func (st *testStream) Close() error {
	return nil
}

func (st *testStream) CloseOnExit() error {
	return nil
}

func makeDummyTestStreams(indexes []int) []sttypes.Stream {
	sts := make([]sttypes.Stream, 0, len(indexes))

	for _, index := range indexes {
		sts = append(sts, &testStream{
			id: makeStreamID(index),
		})
	}
	return sts
}

func makeDummyStreamSets(indexes []int) map[sttypes.StreamID]*stream {
	m := make(map[sttypes.StreamID]*stream)

	for _, index := range indexes {
		st := &testStream{
			id: makeStreamID(index),
		}
		m[st.ID()] = &stream{Stream: st}
	}
	return m
}

func makeStreamID(index int) sttypes.StreamID {
	return sttypes.StreamID(strconv.Itoa(index))
}

type testRequest struct {
	reqID uint64
	index uint64
}

func makeTestRequest(index uint64) *testRequest {
	return &testRequest{
		reqID: 0,
		index: index,
	}
}

func (req *testRequest) ReqID() uint64 {
	return req.reqID
}

func (req *testRequest) SetReqID(rid uint64) {
	req.reqID = rid
}

func (req *testRequest) String() string {
	return fmt.Sprintf("test request %v", req.index)
}

func (req *testRequest) Encode() ([]byte, error) {
	return rlp.EncodeToBytes(struct {
		ReqID uint64
		Index uint64
	}{
		ReqID: req.reqID,
		Index: req.index,
	})
}

func (req *testRequest) checkResponse(rawResp sttypes.Response) error {
	resp, ok := rawResp.(*testResponse)
	if !ok || resp == nil {
		return errors.New("not test Response")
	}
	if req.reqID != resp.reqID {
		return errors.New("request id not expected")
	}
	if req.index != resp.index {
		return errors.New("response id not expected")
	}
	return nil
}

func decodeTestRequest(b []byte) (*testRequest, error) {
	type SerRequest struct {
		ReqID uint64
		Index uint64
	}
	var sr SerRequest
	if err := rlp.DecodeBytes(b, &sr); err != nil {
		return nil, err
	}
	return &testRequest{
		reqID: sr.ReqID,
		index: sr.Index,
	}, nil
}

func (req *testRequest) IsSupportedByProto(spec sttypes.ProtoSpec) bool {
	return true
}

func (req *testRequest) getResponse() *testResponse {
	return &testResponse{
		reqID: req.reqID,
		index: req.index,
	}
}

type testResponse struct {
	reqID uint64
	index uint64
}

func (tr *testResponse) ReqID() uint64 {
	return tr.reqID
}

func (tr *testResponse) String() string {
	return fmt.Sprintf("test response %v", tr.index)
}