github.com/keltia/go-ipfs@v0.3.8-0.20150909044612-210793031c63/metrics/stream/metered_test.go (about)

     1  package meterstream
     2  
     3  import (
     4  	"io"
     5  	"io/ioutil"
     6  	"testing"
     7  
     8  	inet "github.com/ipfs/go-ipfs/p2p/net"
     9  	peer "github.com/ipfs/go-ipfs/p2p/peer"
    10  	protocol "github.com/ipfs/go-ipfs/p2p/protocol"
    11  	u "github.com/ipfs/go-ipfs/util"
    12  )
    13  
    14  type FakeStream struct {
    15  	ReadBuf io.Reader
    16  	inet.Stream
    17  }
    18  
    19  func (fs *FakeStream) Read(b []byte) (int, error) {
    20  	return fs.ReadBuf.Read(b)
    21  }
    22  
    23  func (fs *FakeStream) Write(b []byte) (int, error) {
    24  	return len(b), nil
    25  }
    26  
    27  func TestCallbacksWork(t *testing.T) {
    28  	fake := new(FakeStream)
    29  
    30  	var sent int64
    31  	var recv int64
    32  
    33  	sentCB := func(n int64, proto protocol.ID, p peer.ID) {
    34  		sent += n
    35  	}
    36  
    37  	recvCB := func(n int64, proto protocol.ID, p peer.ID) {
    38  		recv += n
    39  	}
    40  
    41  	ms := newMeteredStream(fake, protocol.ID("TEST"), peer.ID("PEER"), recvCB, sentCB)
    42  
    43  	toWrite := int64(100000)
    44  	toRead := int64(100000)
    45  
    46  	fake.ReadBuf = io.LimitReader(u.NewTimeSeededRand(), toRead)
    47  	writeData := io.LimitReader(u.NewTimeSeededRand(), toWrite)
    48  
    49  	n, err := io.Copy(ms, writeData)
    50  	if err != nil {
    51  		t.Fatal(err)
    52  	}
    53  
    54  	if n != toWrite {
    55  		t.Fatal("incorrect write amount")
    56  	}
    57  
    58  	if toWrite != sent {
    59  		t.Fatal("incorrectly reported writes", toWrite, sent)
    60  	}
    61  
    62  	n, err = io.Copy(ioutil.Discard, ms)
    63  	if err != nil {
    64  		t.Fatal(err)
    65  	}
    66  
    67  	if n != toRead {
    68  		t.Fatal("incorrect read amount")
    69  	}
    70  
    71  	if toRead != recv {
    72  		t.Fatal("incorrectly reported reads")
    73  	}
    74  }