1 files changed, 14 insertions, 6 deletions
diff --git a/src/chunklist_test.go b/src/chunklist_test.go
index a7daa47e..b244ece8 100644
--- a/src/chunklist_test.go
+++ b/src/chunklist_test.go
@@ -11,8 +11,8 @@ func TestChunkList(t *testing.T) {
 	})
 
 	// Snapshot
-	snapshot := cl.Snapshot()
-	if len(snapshot) > 0 {
+	snapshot, count := cl.Snapshot()
+	if len(snapshot) > 0 || count > 0 {
 		t.Error("Snapshot should be empty now")
 	}
 
@@ -26,8 +26,8 @@ func TestChunkList(t *testing.T) {
 	}
 
 	// But the new snapshot should contain the added items
-	snapshot = cl.Snapshot()
-	if len(snapshot) != 1 {
+	snapshot, count = cl.Snapshot()
+	if len(snapshot) != 1 && count != 2 {
 		t.Error("Snapshot should not be empty now")
 	}
 
@@ -55,12 +55,20 @@ func TestChunkList(t *testing.T) {
 	}
 
 	// New snapshot
-	snapshot = cl.Snapshot()
+	snapshot, count = cl.Snapshot()
 	if len(snapshot) != 3 || !snapshot[0].IsFull() ||
-		!snapshot[1].IsFull() || snapshot[2].IsFull() {
+		!snapshot[1].IsFull() || snapshot[2].IsFull() || count != CHUNK_SIZE*2+2 {
 		t.Error("Expected two full chunks and one more chunk")
 	}
 	if len(*snapshot[2]) != 2 {
 		t.Error("Unexpected number of items")
 	}
+
+	cl.Push("hello")
+	cl.Push("world")
+
+	lastChunkCount := len(*snapshot[len(snapshot)-1])
+	if lastChunkCount != 2 {
+		t.Error("Unexpected number of items:", lastChunkCount)
+	}
 }