Mock FlowControlManager using gomock

internal/mocks/flow_control_manager.go

@@ -0,0 +1,140 @@
+// Automatically generated by MockGen. DO NOT EDIT!
+// Source: github.com/lucas-clemente/quic-go/flowcontrol (interfaces: FlowControlManager)
+
+package mocks
+
+import (
+	gomock "github.com/golang/mock/gomock"
+	flowcontrol "github.com/lucas-clemente/quic-go/flowcontrol"
+	protocol "github.com/lucas-clemente/quic-go/protocol"
+)
+
+// Mock of FlowControlManager interface
+type MockFlowControlManager struct {
+	ctrl     *gomock.Controller
+	recorder *_MockFlowControlManagerRecorder
+}
+
+// Recorder for MockFlowControlManager (not exported)
+type _MockFlowControlManagerRecorder struct {
+	mock *MockFlowControlManager
+}
+
+func NewMockFlowControlManager(ctrl *gomock.Controller) *MockFlowControlManager {
+	mock := &MockFlowControlManager{ctrl: ctrl}
+	mock.recorder = &_MockFlowControlManagerRecorder{mock}
+	return mock
+}
+
+func (_m *MockFlowControlManager) EXPECT() *_MockFlowControlManagerRecorder {
+	return _m.recorder
+}
+
+func (_m *MockFlowControlManager) AddBytesRead(_param0 protocol.StreamID, _param1 protocol.ByteCount) error {
+	ret := _m.ctrl.Call(_m, "AddBytesRead", _param0, _param1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+func (_mr *_MockFlowControlManagerRecorder) AddBytesRead(arg0, arg1 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "AddBytesRead", arg0, arg1)
+}
+
+func (_m *MockFlowControlManager) AddBytesSent(_param0 protocol.StreamID, _param1 protocol.ByteCount) error {
+	ret := _m.ctrl.Call(_m, "AddBytesSent", _param0, _param1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+func (_mr *_MockFlowControlManagerRecorder) AddBytesSent(arg0, arg1 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "AddBytesSent", arg0, arg1)
+}
+
+func (_m *MockFlowControlManager) GetReceiveWindow(_param0 protocol.StreamID) (protocol.ByteCount, error) {
+	ret := _m.ctrl.Call(_m, "GetReceiveWindow", _param0)
+	ret0, _ := ret[0].(protocol.ByteCount)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+func (_mr *_MockFlowControlManagerRecorder) GetReceiveWindow(arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "GetReceiveWindow", arg0)
+}
+
+func (_m *MockFlowControlManager) GetWindowUpdates() []flowcontrol.WindowUpdate {
+	ret := _m.ctrl.Call(_m, "GetWindowUpdates")
+	ret0, _ := ret[0].([]flowcontrol.WindowUpdate)
+	return ret0
+}
+
+func (_mr *_MockFlowControlManagerRecorder) GetWindowUpdates() *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "GetWindowUpdates")
+}
+
+func (_m *MockFlowControlManager) NewStream(_param0 protocol.StreamID, _param1 bool) {
+	_m.ctrl.Call(_m, "NewStream", _param0, _param1)
+}
+
+func (_mr *_MockFlowControlManagerRecorder) NewStream(arg0, arg1 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "NewStream", arg0, arg1)
+}
+
+func (_m *MockFlowControlManager) RemainingConnectionWindowSize() protocol.ByteCount {
+	ret := _m.ctrl.Call(_m, "RemainingConnectionWindowSize")
+	ret0, _ := ret[0].(protocol.ByteCount)
+	return ret0
+}
+
+func (_mr *_MockFlowControlManagerRecorder) RemainingConnectionWindowSize() *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "RemainingConnectionWindowSize")
+}
+
+func (_m *MockFlowControlManager) RemoveStream(_param0 protocol.StreamID) {
+	_m.ctrl.Call(_m, "RemoveStream", _param0)
+}
+
+func (_mr *_MockFlowControlManagerRecorder) RemoveStream(arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "RemoveStream", arg0)
+}
+
+func (_m *MockFlowControlManager) ResetStream(_param0 protocol.StreamID, _param1 protocol.ByteCount) error {
+	ret := _m.ctrl.Call(_m, "ResetStream", _param0, _param1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+func (_mr *_MockFlowControlManagerRecorder) ResetStream(arg0, arg1 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "ResetStream", arg0, arg1)
+}
+
+func (_m *MockFlowControlManager) SendWindowSize(_param0 protocol.StreamID) (protocol.ByteCount, error) {
+	ret := _m.ctrl.Call(_m, "SendWindowSize", _param0)
+	ret0, _ := ret[0].(protocol.ByteCount)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+func (_mr *_MockFlowControlManagerRecorder) SendWindowSize(arg0 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "SendWindowSize", arg0)
+}
+
+func (_m *MockFlowControlManager) UpdateHighestReceived(_param0 protocol.StreamID, _param1 protocol.ByteCount) error {
+	ret := _m.ctrl.Call(_m, "UpdateHighestReceived", _param0, _param1)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+func (_mr *_MockFlowControlManagerRecorder) UpdateHighestReceived(arg0, arg1 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "UpdateHighestReceived", arg0, arg1)
+}
+
+func (_m *MockFlowControlManager) UpdateWindow(_param0 protocol.StreamID, _param1 protocol.ByteCount) (bool, error) {
+	ret := _m.ctrl.Call(_m, "UpdateWindow", _param0, _param1)
+	ret0, _ := ret[0].(bool)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+func (_mr *_MockFlowControlManagerRecorder) UpdateWindow(arg0, arg1 interface{}) *gomock.Call {
+	return _mr.mock.ctrl.RecordCall(_mr.mock, "UpdateWindow", arg0, arg1)
+}

internal/mocks/gen.go

@@ -0,0 +1,3 @@
+package mocks
+
+//go:generate mockgen -destination flow_control_manager.go -package mocks github.com/lucas-clemente/quic-go/flowcontrol FlowControlManager

packet_packer_test.go

@@ -48,13 +48,8 @@ var _ = Describe("Packet packer", func() {
 	)
 
 	BeforeEach(func() {
-		fcm := newMockFlowControlHandler()
-		fcm.sendWindowSizes[3] = protocol.MaxByteCount
-		fcm.sendWindowSizes[5] = protocol.MaxByteCount
-		fcm.sendWindowSizes[7] = protocol.MaxByteCount
-
 		cpm := &mockConnectionParametersManager{}
-		streamFramer = newStreamFramer(newStreamsMap(nil, protocol.PerspectiveServer, cpm), fcm)
+		streamFramer = newStreamFramer(newStreamsMap(nil, protocol.PerspectiveServer, cpm), nil)
 
 		packer = &packetPacker{
 			cryptoSetup:           &mockCryptoSetup{encLevelSeal: protocol.EncryptionForwardSecure},

protocol/protocol.go

@@ -31,7 +31,7 @@ type StreamID uint32
 type ByteCount uint64
 
 // MaxByteCount is the maximum value of a ByteCount
-const MaxByteCount = math.MaxUint64
+const MaxByteCount = ByteCount(math.MaxUint64)
 
 // MaxReceivePacketSize maximum packet size of any QUIC packet, based on
 // ethernet's max size, minus the IP and UDP headers. IPv6 has a 40 byte header,

quic_suite_test.go

@@ -1,6 +1,7 @@
 package quic
 
 import (
+	"github.com/golang/mock/gomock"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 
@@ -11,3 +12,13 @@ func TestQuicGo(t *testing.T) {
 	RegisterFailHandler(Fail)
 	RunSpecs(t, "QUIC Suite")
 }
+
+var mockCtrl *gomock.Controller
+
+var _ = BeforeEach(func() {
+	mockCtrl = gomock.NewController(GinkgoT())
+})
+
+var _ = AfterEach(func() {
+	mockCtrl.Finish()
+})

session.go

@@ -31,7 +31,6 @@ type receivedPacket struct {
 var (
 	errRstStreamOnInvalidStream   = errors.New("RST_STREAM received for unknown stream")
 	errWindowUpdateOnClosedStream = errors.New("WINDOW_UPDATE received for an already closed stream")
-	errSessionAlreadyClosed       = errors.New("cannot close session; it was already closed before")
 )
 
 var (

session_test.go

@@ -17,6 +17,7 @@ import (
 	"github.com/lucas-clemente/quic-go/crypto"
 	"github.com/lucas-clemente/quic-go/frames"
 	"github.com/lucas-clemente/quic-go/handshake"
+	"github.com/lucas-clemente/quic-go/internal/mocks"
 	"github.com/lucas-clemente/quic-go/protocol"
 	"github.com/lucas-clemente/quic-go/qerr"
 	"github.com/lucas-clemente/quic-go/testdata"
@@ -474,21 +475,22 @@ var _ = Describe("Session", func() {
 
 		It("passes the byte offset to the flow controller", func() {
 			sess.streamsMap.GetOrOpenStream(5)
-			sess.flowControlManager = newMockFlowControlHandler()
+			fcm := mocks.NewMockFlowControlManager(mockCtrl)
+			sess.flowControlManager = fcm
+			fcm.EXPECT().ResetStream(protocol.StreamID(5), protocol.ByteCount(0x1337))
 			err := sess.handleRstStreamFrame(&frames.RstStreamFrame{
 				StreamID:   5,
 				ByteOffset: 0x1337,
 			})
 			Expect(err).ToNot(HaveOccurred())
-			Expect(sess.flowControlManager.(*mockFlowControlHandler).highestReceivedForStream).To(Equal(protocol.StreamID(5)))
-			Expect(sess.flowControlManager.(*mockFlowControlHandler).highestReceived).To(Equal(protocol.ByteCount(0x1337)))
 		})
 
 		It("returns errors from the flow controller", func() {
-			sess.streamsMap.GetOrOpenStream(5)
-			sess.flowControlManager = newMockFlowControlHandler()
 			testErr := errors.New("flow control violation")
-			sess.flowControlManager.(*mockFlowControlHandler).flowControlViolation = testErr
+			sess.streamsMap.GetOrOpenStream(5)
+			fcm := mocks.NewMockFlowControlManager(mockCtrl)
+			sess.flowControlManager = fcm
+			fcm.EXPECT().ResetStream(protocol.StreamID(5), protocol.ByteCount(0x1337)).Return(testErr)
 			err := sess.handleRstStreamFrame(&frames.RstStreamFrame{
 				StreamID:   5,
 				ByteOffset: 0x1337,
@@ -1062,9 +1064,11 @@ var _ = Describe("Session", func() {
 			It("retransmits a WindowUpdates if it hasn't already sent a WindowUpdate with a higher ByteOffset", func() {
 				_, err := sess.GetOrOpenStream(5)
 				Expect(err).ToNot(HaveOccurred())
-				fc := newMockFlowControlHandler()
-				fc.receiveWindow = 0x1000
-				sess.flowControlManager = fc
+				fcm := mocks.NewMockFlowControlManager(mockCtrl)
+				sess.flowControlManager = fcm
+				fcm.EXPECT().GetWindowUpdates()
+				fcm.EXPECT().GetReceiveWindow(protocol.StreamID(5)).Return(protocol.ByteCount(0x1000), nil)
+				fcm.EXPECT().GetWindowUpdates()
 				wuf := &frames.WindowUpdateFrame{
 					StreamID:   5,
 					ByteOffset: 0x1000,
@@ -1082,9 +1086,10 @@ var _ = Describe("Session", func() {
 			It("doesn't retransmit WindowUpdates if it already sent a WindowUpdate with a higher ByteOffset", func() {
 				_, err := sess.GetOrOpenStream(5)
 				Expect(err).ToNot(HaveOccurred())
-				fc := newMockFlowControlHandler()
-				fc.receiveWindow = 0x2000
-				sess.flowControlManager = fc
+				fcm := mocks.NewMockFlowControlManager(mockCtrl)
+				sess.flowControlManager = fcm
+				fcm.EXPECT().GetWindowUpdates()
+				fcm.EXPECT().GetReceiveWindow(protocol.StreamID(5)).Return(protocol.ByteCount(0x2000), nil)
 				sph.retransmissionQueue = []*ackhandler.Packet{{
 					Frames: []frames.Frame{&frames.WindowUpdateFrame{
 						StreamID:   5,