From 1d7aa22424b7f4398e7f640b22367f88032c86c5 Mon Sep 17 00:00:00 2001
From: Vlad <gluk256@gmail.com>
Date: Wed, 6 Mar 2019 15:47:31 +0400
Subject: [PATCH] swarm/pss: WaitTillSnapshotLoaded() fixed

---
 swarm/network/simulation/kademlia.go | 39 ++++++++++++++++++++++++----
 swarm/pss/snapshot_test.go           |  3 ++-
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/swarm/network/simulation/kademlia.go b/swarm/network/simulation/kademlia.go
index 1839741119..e9d34790ce 100644
--- a/swarm/network/simulation/kademlia.go
+++ b/swarm/network/simulation/kademlia.go
@@ -101,7 +101,7 @@ func (s *Simulation) kademlias() (ks map[enode.ID]*network.Kademlia) {
 
 func (s *Simulation) WaitTillSnapshotRecreated(ctx context.Context, snap simulations.Snapshot) error {
 	expected := listSnapshotConnections(snap.Conns)
-	ticker := time.NewTicker(256 * time.Millisecond) // todo: reduce
+	ticker := time.NewTicker(16 * time.Millisecond)
 	defer ticker.Stop()
 
 	for {
@@ -125,6 +125,9 @@ func (s *Simulation) listActualConnections() (res []uint64) {
 			return true
 		})
 	}
+
+	// only list those connections that appear twice (both peers should recognize connection as active)
+	res = removeDuplicatesAndSingletons(res)
 	return res
 }
 
@@ -154,15 +157,41 @@ func isAllDeployed(expected []uint64, actual []uint64) bool {
 			// remove value c from exp
 			for i := 0; i < len(exp); i++ {
 				if exp[i] == c {
-					last := len(exp) - 1
-					if last == 0 {
+					exp = removeListElement(exp, i)
+					if len(exp) == 0 {
 						return true
 					}
-					exp[i] = exp[last]
-					exp = exp[:last]
 				}
 			}
 		}
 	}
 	return len(exp) == 0
 }
+
+func removeListElement(arr []uint64, i int) []uint64 {
+	last := len(arr)-1
+	arr[i] = arr[last]
+	arr = arr[:last]
+	return arr
+}
+
+func removeDuplicatesAndSingletons(arr []uint64) []uint64 {
+	for i := 0; i < len(arr); {
+		found := false
+		for j := i+1; j < len(arr); j++ {
+			if arr[i] == arr[j] {
+				arr = removeListElement(arr, j) // remove duplicate
+				found = true
+				break
+			}
+		}
+
+		if found {
+			i++
+		} else {
+			arr = removeListElement(arr, i) // remove singleton
+		}
+	}
+
+	return arr
+}
diff --git a/swarm/pss/snapshot_test.go b/swarm/pss/snapshot_test.go
index b0d67da270..abf0823a24 100644
--- a/swarm/pss/snapshot_test.go
+++ b/swarm/pss/snapshot_test.go
@@ -76,6 +76,7 @@ func resetTestVariables() {
 	maxMessages = 0
 	msgCnt = 0
 	msgs = nil
+	sim = nil
 
 	kademlias     = make(map[enode.ID]*network.Kademlia)
 	nodeAddrs     = make(map[enode.ID][]byte)
@@ -165,7 +166,7 @@ func assingTestVariables(sim *simulation.Simulation, msgCount int) {
 
 			if po >= depth {
 				maxMessages++
-				allowed[i] = append(recipients[i], nod.ID())
+				allowed[i] = append(allowed[i], nod.ID())
 				allowedMsgs[nod.ID()] = append(allowedMsgs[nod.ID()], uint64(i))
 			}