From 413c398cfff86e8420de1c9f91634d55a091efec Mon Sep 17 00:00:00 2001
From: Diogo Behrens <diogo.behrens@lovoo.com>
Date: Mon, 9 Apr 2018 16:02:51 +0200
Subject: [PATCH] fixing multi-iterator test cases

---
 storage/memory.go              |  3 +-
 storage/multi_iterator.go      |  9 +++---
 storage/multi_iterator_test.go | 50 ++++++++++++++++++++--------------
 3 files changed, 36 insertions(+), 26 deletions(-)

diff --git a/storage/memory.go b/storage/memory.go
index f60e440b..1e9c6d70 100644
--- a/storage/memory.go
+++ b/storage/memory.go
@@ -3,6 +3,7 @@ package storage
 import (
 	"bytes"
 	"fmt"
+	"strings"
 
 	"github.com/syndtr/goleveldb/leveldb/util"
 )
@@ -50,7 +51,7 @@ func (i *memiter) Seek(key []byte) bool {
 	seek := make(map[string][]byte)
 	keys := []string{}
 	for k, v := range i.storage {
-		if bytes.ContainsAny(key, k) {
+		if strings.Contains(k, string(key)) {
 			keys = append(keys, k)
 			seek[k] = v
 		}
diff --git a/storage/multi_iterator.go b/storage/multi_iterator.go
index 9708ed5d..2fcdaf11 100644
--- a/storage/multi_iterator.go
+++ b/storage/multi_iterator.go
@@ -14,14 +14,13 @@ func NewMultiIterator(iters []Iterator) Iterator {
 	return &multiIterator{current: 0, iters: iters}
 }
 
-func (m *multiIterator) Next() (next bool) {
+func (m *multiIterator) Next() bool {
 	for ; m.current < len(m.iters); m.current++ {
-		next = m.iters[m.current].Next()
-		if next {
-			break
+		if m.iters[m.current].Next() {
+			return true
 		}
 	}
-	return next
+	return false
 }
 
 func (m *multiIterator) Key() []byte {
diff --git a/storage/multi_iterator_test.go b/storage/multi_iterator_test.go
index a32d14e4..13723f6f 100644
--- a/storage/multi_iterator_test.go
+++ b/storage/multi_iterator_test.go
@@ -8,16 +8,17 @@ import (
 )
 
 func TestMultiIterator(t *testing.T) {
-	numStorages := 3
-	numValues := 3
-
-	storages := make([]Storage, numStorages)
-	expected := map[string]string{}
+	var (
+		numStorages = 3
+		numValues   = 3
+		storages    = make([]Storage, numStorages)
+		expected    = make(map[string]string)
+	)
 
 	for i := 0; i < numStorages; i++ {
 		storages[i] = NewMemory()
 		for j := 0; j < numValues; j++ {
-			key := fmt.Sprintf("storage-%d", i)
+			key := fmt.Sprintf("key-%d-%d", i, j)
 			val := fmt.Sprintf("value-%d", j)
 			expected[key] = val
 			storages[i].Set(key, []byte(val))
@@ -41,13 +42,11 @@ func TestMultiIterator(t *testing.T) {
 	}
 	ensure.DeepEqual(t, count, len(expected))
 
-	k := []byte("storage-0")
+	k := []byte("key-2")
 	iter = NewMultiIterator(iters)
 	ensure.True(t, iter.Seek(k), "seek return false should return true")
-	ensure.True(t, iter.Next(), "Iterator should have a value")
-	ensure.DeepEqual(t, iter.Key(), k, "key mismatch")
 
-	total := 1
+	total := 0
 	for iter.Next() {
 		_, err := iter.Value()
 		ensure.Nil(t, err)
@@ -85,18 +84,21 @@ func TestMultiIteratorOneValue(t *testing.T) {
 }
 
 func TestMultiIteratorMixedValues(t *testing.T) {
-	storages := make([]Storage, 3)
-	var expected []string
-	n := 0
+	var (
+		n        = 0
+		storages = make([]Storage, 3)
+		expected = make(map[string]string)
+		counts   = make(map[string]int)
+	)
 
 	// first storage has two values
 	storages[0] = NewMemory()
 	for i := 0; i < 2; i++ {
 		key := fmt.Sprintf("key-%d", n)
 		val := fmt.Sprintf("value-%d", n)
-		n++
-		expected = append(expected, val)
+		expected[key] = val
 		storages[0].Set(key, []byte(val))
+		n++
 	}
 
 	// second storage has 0 values
@@ -107,9 +109,9 @@ func TestMultiIteratorMixedValues(t *testing.T) {
 	for i := 0; i < 3; i++ {
 		key := fmt.Sprintf("key-%d", n)
 		val := fmt.Sprintf("value-%d", n)
-		n++
-		expected = append(expected, val)
+		expected[key] = val
 		storages[2].Set(key, []byte(val))
+		n++
 	}
 
 	iters := make([]Iterator, len(storages))
@@ -120,12 +122,20 @@ func TestMultiIteratorMixedValues(t *testing.T) {
 	}
 
 	iter := NewMultiIterator(iters)
-	count := 0
 	for iter.Next() {
 		val, err := iter.Value()
 		ensure.Nil(t, err)
-		ensure.DeepEqual(t, expected[count], string(val))
-		count++
+		ensure.DeepEqual(t, expected[string(iter.Key())], string(val))
+		counts[string(iter.Key())]++
+	}
+
+	// uniques and totals should match
+	unique := 0
+	count := 0
+	for _, c := range counts {
+		unique++
+		count += c
 	}
 	ensure.DeepEqual(t, count, len(expected))
+	ensure.DeepEqual(t, unique, len(expected))
 }