Merge pull request #2293 from lucas-clemente/mock-connection

use a mock connection in tests

closed_session_test.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"time"
 
+	"github.com/golang/mock/gomock"
 	"github.com/lucas-clemente/quic-go/internal/protocol"
 	"github.com/lucas-clemente/quic-go/internal/utils"
 
@@ -14,11 +15,11 @@ import (
 var _ = Describe("Closed local session", func() {
 	var (
 		sess  packetHandler
-		mconn *mockConnection
+		mconn *MockConnection
 	)
 
 	BeforeEach(func() {
-		mconn = newMockConnection()
+		mconn = NewMockConnection(mockCtrl)
 		sess = newClosedLocalSession(mconn, []byte("close"), protocol.PerspectiveClient, utils.DefaultLogger)
 	})
 
@@ -33,12 +34,14 @@ var _ = Describe("Closed local session", func() {
 	})
 
 	It("repeats the packet containing the CONNECTION_CLOSE frame", func() {
+		written := make(chan []byte)
+		mconn.EXPECT().Write(gomock.Any()).Do(func(p []byte) { written <- p }).AnyTimes()
 		for i := 1; i <= 20; i++ {
 			sess.handlePacket(&receivedPacket{})
 			if i == 1 || i == 2 || i == 4 || i == 8 || i == 16 {
-				Eventually(mconn.written).Should(Receive(Equal([]byte("close")))) // receive the CONNECTION_CLOSE
+				Eventually(written).Should(Receive(Equal([]byte("close")))) // receive the CONNECTION_CLOSE
 			} else {
-				Consistently(mconn.written, 10*time.Millisecond).Should(HaveLen(0))
+				Consistently(written, 10*time.Millisecond).Should(HaveLen(0))
 			}
 		}
 		// stop the session

mockgen.go

@@ -1,5 +1,6 @@
 package quic
 
+//go:generate sh -c "./mockgen_private.sh quic mock_connection_test.go github.com/lucas-clemente/quic-go connection"
 //go:generate sh -c "./mockgen_private.sh quic mock_stream_internal_test.go github.com/lucas-clemente/quic-go streamI"
 //go:generate sh -c "./mockgen_private.sh quic mock_crypto_stream_test.go github.com/lucas-clemente/quic-go cryptoStream"
 //go:generate sh -c "./mockgen_private.sh quic mock_receive_stream_internal_test.go github.com/lucas-clemente/quic-go receiveStreamI"

send_queue_test.go

@@ -1,16 +1,17 @@
 package quic
 
 import (
+	"github.com/golang/mock/gomock"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 )
 
 var _ = Describe("Send Queue", func() {
 	var q *sendQueue
-	var c *mockConnection
+	var c *MockConnection
 
 	BeforeEach(func() {
-		c = newMockConnection()
+		c = NewMockConnection(mockCtrl)
 		q = newSendQueue(c)
 	})
 
@@ -25,23 +26,29 @@ var _ = Describe("Send Queue", func() {
 	}
 
 	It("sends a packet", func() {
-		q.Send(getPacket([]byte("foobar")))
+		p := getPacket([]byte("foobar"))
+		q.Send(p)
 
+		written := make(chan struct{})
+		c.EXPECT().Write(p.raw).Do(func([]byte) { close(written) })
 		done := make(chan struct{})
 		go func() {
 			defer GinkgoRecover()
 			q.Run()
 			close(done)
 		}()
 
-		Eventually(c.written).Should(Receive(Equal([]byte("foobar"))))
+		Eventually(written).Should(BeClosed())
 		q.Close()
 		Eventually(done).Should(BeClosed())
 	})
 
 	It("blocks sending when too many packets are queued", func() {
 		q.Send(getPacket([]byte("foobar")))
 
+		written := make(chan []byte, 2)
+		c.EXPECT().Write(gomock.Any()).Do(func(p []byte) { written <- p }).Times(2)
+
 		sent := make(chan struct{})
 		go func() {
 			defer GinkgoRecover()
@@ -58,8 +65,8 @@ var _ = Describe("Send Queue", func() {
 			close(done)
 		}()
 
-		Eventually(c.written).Should(Receive(Equal([]byte("foobar"))))
-		Eventually(c.written).Should(Receive(Equal([]byte("raboof"))))
+		Eventually(written).Should(Receive(Equal([]byte("foobar"))))
+		Eventually(written).Should(Receive(Equal([]byte("raboof"))))
 		q.Close()
 		Eventually(done).Should(BeClosed())
 	})