From 5270b807526091430b798ba82582e2e922e5bf7a Mon Sep 17 00:00:00 2001
From: Attila Gazso <agazso@gmail.com>
Date: Tue, 31 Jul 2018 12:15:52 +0200
Subject: [PATCH] swarm/network: Changed LightNode handhsake test to table
 driven

---
 swarm/network/protocol_test.go | 58 ++++++++++++++--------------------
 1 file changed, 24 insertions(+), 34 deletions(-)

diff --git a/swarm/network/protocol_test.go b/swarm/network/protocol_test.go
index a6f26e5f02..2578430292 100644
--- a/swarm/network/protocol_test.go
+++ b/swarm/network/protocol_test.go
@@ -255,44 +255,34 @@ func TestBzzHandshakeSuccess(t *testing.T) {
 	}
 }
 
-func TestBzzHandshakeLightNodeOff(t *testing.T) {
-	randomAddr := RandomAddr()
-	pt := newBzzHandshakeTester(t, 1, randomAddr, false)
-	id := pt.IDs[0]
-	addr := NewAddrFromNodeID(id)
-	peerLightNode := false
-
-	err := pt.testHandshake(
-		correctBzzHandshake(randomAddr, false),
-		&HandshakeMsg{Version: TestProtocolVersion, NetworkID: TestProtocolNetworkID, Addr: addr, LightNode: peerLightNode},
-	)
-
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	if pt.bzz.handshakes[id].LightNode != peerLightNode {
-		t.Fatalf("peer LightNode flag is %v, should be %v", pt.bzz.handshakes[id].LightNode, peerLightNode)
-	}
+var lightNodeTests = []struct {
+	name      string
+	lightNode bool
+}{
+	{"on", true},
+	{"off", false},
 }
 
-func TestBzzHandshakeLightNodeOn(t *testing.T) {
-	randomAddr := RandomAddr()
-	pt := newBzzHandshakeTester(t, 1, randomAddr, false)
-	id := pt.IDs[0]
-	addr := NewAddrFromNodeID(id)
-	peerLightNode := true
+func TestBzzHandshakeLightNode(t *testing.T) {
+	for _, test := range lightNodeTests {
+		t.Run(test.name, func(t *testing.T) {
+			randomAddr := RandomAddr()
+			pt := newBzzHandshakeTester(t, 1, randomAddr, false)
+			id := pt.IDs[0]
+			addr := NewAddrFromNodeID(id)
 
-	err := pt.testHandshake(
-		correctBzzHandshake(randomAddr, false),
-		&HandshakeMsg{Version: TestProtocolVersion, NetworkID: TestProtocolNetworkID, Addr: addr, LightNode: peerLightNode},
-	)
+			err := pt.testHandshake(
+				correctBzzHandshake(randomAddr, false),
+				&HandshakeMsg{Version: TestProtocolVersion, NetworkID: TestProtocolNetworkID, Addr: addr, LightNode: test.lightNode},
+			)
 
-	if err != nil {
-		t.Fatal(err)
-	}
+			if err != nil {
+				t.Fatal(err)
+			}
 
-	if pt.bzz.handshakes[id].LightNode != peerLightNode {
-		t.Fatalf("peer LightNode flag is %v, should be %v", pt.bzz.handshakes[id].LightNode, peerLightNode)
+			if pt.bzz.handshakes[id].LightNode != test.lightNode {
+				t.Fatalf("peer LightNode flag is %v, should be %v", pt.bzz.handshakes[id].LightNode, test.lightNode)
+			}
+		})
 	}
 }