From 395b0a312dd078b115cfdeaa0888a8cf9afa0e95 Mon Sep 17 00:00:00 2001
From: Sreekanth Sivasankaran <sreekanth.sivasankaran@couchbase.com>
Date: Mon, 5 Mar 2018 17:02:58 +0530
Subject: [PATCH] adding UTs

---
 index/scorch/segment/zap/merge.go      |  4 +-
 index/scorch/segment/zap/merge_test.go | 96 ++++++++++++++++++++++++--
 2 files changed, 95 insertions(+), 5 deletions(-)

diff --git a/index/scorch/segment/zap/merge.go b/index/scorch/segment/zap/merge.go
index a9383f20..9399046b 100644
--- a/index/scorch/segment/zap/merge.go
+++ b/index/scorch/segment/zap/merge.go
@@ -97,7 +97,9 @@ func Merge(segments []*Segment, drops []*roaring.Bitmap, path string,
 		return nil, err
 	}
 
-	stats.ReportBytesWritten(uint64(cr.Count()))
+	if stats != nil {
+		stats.ReportBytesWritten(uint64(cr.Count()))
+	}
 
 	return newDocNums, nil
 }
diff --git a/index/scorch/segment/zap/merge_test.go b/index/scorch/segment/zap/merge_test.go
index bb09f831..2ba0b373 100644
--- a/index/scorch/segment/zap/merge_test.go
+++ b/index/scorch/segment/zap/merge_test.go
@@ -20,6 +20,7 @@ import (
 	"reflect"
 	"sort"
 	"strings"
+	"sync/atomic"
 	"testing"
 
 	"github.com/RoaringBitmap/roaring"
@@ -72,7 +73,7 @@ func TestMerge(t *testing.T) {
 	segsToMerge[0] = segment.(*Segment)
 	segsToMerge[1] = segment2.(*Segment)
 
-	_, err = Merge(segsToMerge, []*roaring.Bitmap{nil, nil}, "/tmp/scorch3.zap", 1024)
+	_, err = Merge(segsToMerge, []*roaring.Bitmap{nil, nil}, "/tmp/scorch3.zap", 1024, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -176,7 +177,7 @@ func testMergeWithEmptySegments(t *testing.T, before bool, numEmptySegments int)
 
 	drops := make([]*roaring.Bitmap, len(segsToMerge))
 
-	_, err = Merge(segsToMerge, drops, "/tmp/scorch3.zap", 1024)
+	_, err = Merge(segsToMerge, drops, "/tmp/scorch3.zap", 1024, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -218,7 +219,7 @@ func testMergeWithSelf(t *testing.T, segCur *Segment, expectedCount uint64) {
 		segsToMerge := make([]*Segment, 1)
 		segsToMerge[0] = segCur
 
-		_, err := Merge(segsToMerge, []*roaring.Bitmap{nil, nil}, "/tmp/"+fname, 1024)
+		_, err := Merge(segsToMerge, []*roaring.Bitmap{nil, nil}, "/tmp/"+fname, 1024, nil)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -590,7 +591,7 @@ func testMergeWithUpdates(t *testing.T, segmentDocIds [][]string, docsToDrop []*
 func testMergeAndDropSegments(t *testing.T, segsToMerge []*Segment, docsToDrop []*roaring.Bitmap, expectedNumDocs uint64) {
 	_ = os.RemoveAll("/tmp/scorch-merged.zap")
 
-	_, err := Merge(segsToMerge, docsToDrop, "/tmp/scorch-merged.zap", 1024)
+	_, err := Merge(segsToMerge, docsToDrop, "/tmp/scorch-merged.zap", 1024, nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -782,3 +783,90 @@ func buildMemSegmentMultiHelper(docIds []string) *mem.Segment {
 
 	return segment
 }
+
+type statTest struct {
+	totalWrittenBytes uint64
+}
+
+func (s *statTest) ReportBytesWritten(numBytesWritten uint64) {
+	atomic.AddUint64(&s.totalWrittenBytes, numBytesWritten)
+}
+
+func TestMergeBytesWritten(t *testing.T) {
+	_ = os.RemoveAll("/tmp/scorch.zap")
+	_ = os.RemoveAll("/tmp/scorch2.zap")
+	_ = os.RemoveAll("/tmp/scorch3.zap")
+
+	memSegment := buildMemSegmentMulti()
+	err := PersistSegment(memSegment, "/tmp/scorch.zap", 1024)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	memSegment2 := buildMemSegmentMulti2()
+	err = PersistSegment(memSegment2, "/tmp/scorch2.zap", 1024)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	segment, err := Open("/tmp/scorch.zap")
+	if err != nil {
+		t.Fatalf("error opening segment: %v", err)
+	}
+	defer func() {
+		cerr := segment.Close()
+		if cerr != nil {
+			t.Fatalf("error closing segment: %v", err)
+		}
+	}()
+
+	segment2, err := Open("/tmp/scorch2.zap")
+	if err != nil {
+		t.Fatalf("error opening segment: %v", err)
+	}
+	defer func() {
+		cerr := segment2.Close()
+		if cerr != nil {
+			t.Fatalf("error closing segment: %v", err)
+		}
+	}()
+
+	segsToMerge := make([]*Segment, 2)
+	segsToMerge[0] = segment.(*Segment)
+	segsToMerge[1] = segment2.(*Segment)
+
+	reporter := &statTest{}
+
+	_, err = Merge(segsToMerge, []*roaring.Bitmap{nil, nil}, "/tmp/scorch3.zap", 1024, reporter)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if reporter.totalWrittenBytes == 0 {
+		t.Fatalf("expected a non zero total_compaction_written_bytes")
+	}
+
+	segm, err := Open("/tmp/scorch3.zap")
+	if err != nil {
+		t.Fatalf("error opening merged segment: %v", err)
+	}
+	seg3 := segm.(*Segment)
+	defer func() {
+		cerr := seg3.Close()
+		if cerr != nil {
+			t.Fatalf("error closing segment: %v", err)
+		}
+	}()
+
+	if seg3.Path() != "/tmp/scorch3.zap" {
+		t.Fatalf("wrong path")
+	}
+	if seg3.Count() != 4 {
+		t.Fatalf("wrong count")
+	}
+	if len(seg3.Fields()) != 5 {
+		t.Fatalf("wrong # fields: %#v\n", seg3.Fields())
+	}
+
+	testMergeWithSelf(t, seg3, 4)
+}