From a731c71ed65db9867237d7ab576cf0afef9d188f Mon Sep 17 00:00:00 2001
From: Jake Hillion <jake@hillion.co.uk>
Date: Mon, 2 Nov 2020 17:29:02 +0000
Subject: [PATCH] fixed tests

---
 mocks/conn.go    | 6 ++++++
 tcp/flow.go      | 2 +-
 tcp/flow_test.go | 6 +++---
 3 files changed, 10 insertions(+), 4 deletions(-)

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)