Commit f34a13b3 authored by Janos Guljas's avatar Janos Guljas

update p2p/mock package

parent 44adb6ec
...@@ -9,47 +9,55 @@ import ( ...@@ -9,47 +9,55 @@ import (
ma "github.com/multiformats/go-multiaddr" ma "github.com/multiformats/go-multiaddr"
) )
type Streamer struct { type Recorder struct {
In, Out *Recorder in, out *record
handler func(p2p.Peer) handler func(p2p.Peer)
} }
func NewStreamer(handler func(p2p.Peer)) *Streamer { func NewRecorder(handler func(p2p.Peer)) *Recorder {
return &Streamer{ return &Recorder{
In: NewRecorder(), in: newRecord(),
Out: NewRecorder(), out: newRecord(),
handler: handler, handler: handler,
} }
} }
func (s *Streamer) NewStream(ctx context.Context, peerID, protocolName, streamName, version string) (p2p.Stream, error) { func (r *Recorder) NewStream(ctx context.Context, peerID, protocolName, streamName, version string) (p2p.Stream, error) {
out := NewStream(s.In, s.Out) out := newStream(r.in, r.out)
in := NewStream(s.Out, s.In) in := newStream(r.out, r.in)
go s.handler(p2p.Peer{ go r.handler(p2p.Peer{
Addr: ma.StringCast(peerID), Addr: ma.StringCast(peerID),
Stream: in, Stream: in,
}) })
return out, nil return out, nil
} }
type Stream struct { func (r *Recorder) In() []byte {
return r.in.bytes()
}
func (r *Recorder) Out() []byte {
return r.out.bytes()
}
type stream struct {
in io.WriteCloser in io.WriteCloser
out io.ReadCloser out io.ReadCloser
} }
func NewStream(in io.WriteCloser, out io.ReadCloser) *Stream { func newStream(in io.WriteCloser, out io.ReadCloser) *stream {
return &Stream{in: in, out: out} return &stream{in: in, out: out}
} }
func (s *Stream) Read(p []byte) (int, error) { func (s *stream) Read(p []byte) (int, error) {
return s.out.Read(p) return s.out.Read(p)
} }
func (s *Stream) Write(p []byte) (int, error) { func (s *stream) Write(p []byte) (int, error) {
return s.in.Write(p) return s.in.Write(p)
} }
func (s *Stream) Close() error { func (s *stream) Close() error {
if err := s.in.Close(); err != nil { if err := s.in.Close(); err != nil {
return err return err
} }
...@@ -59,20 +67,20 @@ func (s *Stream) Close() error { ...@@ -59,20 +67,20 @@ func (s *Stream) Close() error {
return nil return nil
} }
type Recorder struct { type record struct {
b []byte b []byte
c int c int
closed bool closed bool
cond *sync.Cond cond *sync.Cond
} }
func NewRecorder() *Recorder { func newRecord() *record {
return &Recorder{ return &record{
cond: sync.NewCond(new(sync.Mutex)), cond: sync.NewCond(new(sync.Mutex)),
} }
} }
func (r *Recorder) Read(p []byte) (n int, err error) { func (r *record) Read(p []byte) (n int, err error) {
r.cond.L.Lock() r.cond.L.Lock()
defer r.cond.L.Unlock() defer r.cond.L.Unlock()
...@@ -91,7 +99,7 @@ func (r *Recorder) Read(p []byte) (n int, err error) { ...@@ -91,7 +99,7 @@ func (r *Recorder) Read(p []byte) (n int, err error) {
return n, err return n, err
} }
func (r *Recorder) Write(p []byte) (int, error) { func (r *record) Write(p []byte) (int, error) {
r.cond.L.Lock() r.cond.L.Lock()
defer r.cond.L.Unlock() defer r.cond.L.Unlock()
...@@ -101,7 +109,7 @@ func (r *Recorder) Write(p []byte) (int, error) { ...@@ -101,7 +109,7 @@ func (r *Recorder) Write(p []byte) (int, error) {
return len(p), nil return len(p), nil
} }
func (r *Recorder) Close() error { func (r *record) Close() error {
r.cond.L.Lock() r.cond.L.Lock()
defer r.cond.L.Unlock() defer r.cond.L.Unlock()
...@@ -111,6 +119,6 @@ func (r *Recorder) Close() error { ...@@ -111,6 +119,6 @@ func (r *Recorder) Close() error {
return nil return nil
} }
func (r *Recorder) Bytes() []byte { func (r *record) bytes() []byte {
return r.b return r.b
} }
...@@ -15,11 +15,11 @@ func TestPing(t *testing.T) { ...@@ -15,11 +15,11 @@ func TestPing(t *testing.T) {
// create a pingpong server that handles the incoming stream // create a pingpong server that handles the incoming stream
server := pingpong.New(nil) server := pingpong.New(nil)
// setup the mock streamer to record stream data // setup the stream recorder to record stream data
streamer := mock.NewStreamer(server.Handler) recorder := mock.NewRecorder(server.Handler)
// create a pingpong client that will do pinging // create a pingpong client that will do pinging
client := pingpong.New(streamer) client := pingpong.New(recorder)
// ping // ping
greetings := []string{"hey", "there", "fella"} greetings := []string{"hey", "there", "fella"}
...@@ -36,7 +36,7 @@ func TestPing(t *testing.T) { ...@@ -36,7 +36,7 @@ func TestPing(t *testing.T) {
// validate received ping greetings from the client // validate received ping greetings from the client
wantGreetings := greetings wantGreetings := greetings
messages, err := protobuf.ReadMessages( messages, err := protobuf.ReadMessages(
bytes.NewReader(streamer.In.Bytes()), bytes.NewReader(recorder.In()),
func() protobuf.Message { return new(pingpong.Ping) }, func() protobuf.Message { return new(pingpong.Ping) },
) )
if err != nil { if err != nil {
...@@ -56,7 +56,7 @@ func TestPing(t *testing.T) { ...@@ -56,7 +56,7 @@ func TestPing(t *testing.T) {
wantResponses = append(wantResponses, "{"+g+"}") wantResponses = append(wantResponses, "{"+g+"}")
} }
messages, err = protobuf.ReadMessages( messages, err = protobuf.ReadMessages(
bytes.NewReader(streamer.Out.Bytes()), bytes.NewReader(recorder.Out()),
func() protobuf.Message { return new(pingpong.Pong) }, func() protobuf.Message { return new(pingpong.Pong) },
) )
if err != nil { if err != nil {
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment