diff --git a/mocks/conn.go b/mocks/conn.go index b817911..3ad384f 100644 --- a/mocks/conn.go +++ b/mocks/conn.go @@ -1,5 +1,7 @@ package mocks +import "time" + type MockPerfectBiConn struct { directionA chan byte directionB chan byte @@ -25,6 +27,10 @@ type MockPerfectConn struct { outbound chan byte } +func (c MockPerfectConn) SetWriteDeadline(time.Time) error { + return nil +} + func (c MockPerfectConn) Read(p []byte) (n int, err error) { for i := range p { if i == 0 { diff --git a/tcp/flow.go b/tcp/flow.go index 48db06f..31cf8c9 100644 --- a/tcp/flow.go +++ b/tcp/flow.go @@ -99,7 +99,7 @@ func (f *Flow) consumeMarshalled(data []byte) error { binary.LittleEndian.PutUint32(prefixedData, uint32(len(data))) copy(prefixedData[4:], data) - err := f.conn.SetWriteDeadline(time.Now().Add(5*time.Second)) + err := f.conn.SetWriteDeadline(time.Now().Add(5 * time.Second)) if err != nil { return err } diff --git a/tcp/flow_test.go b/tcp/flow_test.go index dace021..88fbf1c 100644 --- a/tcp/flow_test.go +++ b/tcp/flow_test.go @@ -17,7 +17,7 @@ func TestFlow_Consume(t *testing.T) { t.Run("Length", func(t *testing.T) { testConn := mocks.NewMockPerfectBiConn(100) - flowA := Flow{conn: testConn.SideA()} + flowA := Flow{conn: testConn.SideA(), isAlive: true} err := flowA.Consume(testPacket, testMac) require.Nil(t, err) @@ -41,7 +41,7 @@ func TestFlow_Produce(t *testing.T) { t.Run("Length", func(t *testing.T) { testConn := mocks.NewMockPerfectBiConn(100) - flowA := Flow{conn: testConn.SideA()} + flowA := Flow{conn: testConn.SideA(), isAlive: true} _, err := testConn.SideB().Write(testMarshalled) require.Nil(t, err) @@ -54,7 +54,7 @@ func TestFlow_Produce(t *testing.T) { t.Run("Value", func(t *testing.T) { testConn := mocks.NewMockPerfectBiConn(100) - flowA := Flow{conn: testConn.SideA()} + flowA := Flow{conn: testConn.SideA(), isAlive: true} _, err := testConn.SideB().Write(testMarshalled) require.Nil(t, err)