Fix stats test

Author: buraksezer

client/client_test.go

@@ -623,6 +623,8 @@ func TestClient_Stats(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Expected nil. Got: %v", err)
 	}
+
+	var totalByKeyCount int
 	var total int
 	for partID, part := range s.Partitions {
 		total += part.TotalKeyCount
@@ -633,13 +635,17 @@ func TestClient_Stats(t *testing.T) {
 			t.Fatalf("Expected PreviosOwners list is empty. "+
 				"Got: %v for PartID: %d", part.PreviousOwners, partID)
 		}
-		if part.KeyCount != 1 {
-			t.Fatalf("Expected DMap count: 1. Got: %d", part.KeyCount)
+		if part.KeyCount <= 0 {
+			t.Fatalf("Expected KeyCount is bigger than 0. Got: %d", part.KeyCount)
 		}
+		totalByKeyCount += part.KeyCount
 		if part.Owner.String() != addr {
 			t.Fatalf("Expected partition owner: %s. Got: %s", addr, part.Owner)
 		}
 	}
+	if totalByKeyCount != 100 {
+		t.Fatalf("Expected total key count in stats is 100. Got: %d", total)
+	}
 	if total != 100 {
 		t.Fatalf("Expected total key count in stats is 100. Got: %d", total)
 	}