1 files changed, 138 insertions, 0 deletions
diff --git a/warc/warc_test.go b/warc/warc_test.go
new file mode 100644
index 0000000..5e2de2f
--- /dev/null
+++ b/warc/warc_test.go
@@ -0,0 +1,138 @@
+package warc
+
+import (
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"sync"
+	"testing"
+)
+
+var testData = []byte("this is some very interesting test data of non-zero size")
+
+func writeRecords(w *Writer, n int) error {
+	for i := 0; i < n; i++ {
+		hdr := NewHeader()
+		rec, err := w.NewRecord(hdr)
+		if err != nil {
+			return fmt.Errorf("NewRecord: %v", err)
+		}
+		_, err = rec.Write(testData)
+		if err != nil {
+			return fmt.Errorf("record Write: %v", err)
+		}
+		if err := rec.Close(); err != nil {
+			return fmt.Errorf("record Close: %v", err)
+		}
+	}
+	return nil
+}
+
+func writeManyRecords(t testing.TB, w *Writer, n int) {
+	if err := writeRecords(w, n); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func writeManyRecordsConcurrently(t testing.TB, w *Writer, n, nproc int) {
+	startCh := make(chan struct{})
+	errCh := make(chan error, nproc+1)
+	var wg sync.WaitGroup
+
+	for i := 0; i < nproc; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			<-startCh
+			if err := writeRecords(w, n); err != nil {
+				errCh <- err
+			}
+		}()
+	}
+	go func() {
+		wg.Wait()
+		errCh <- nil
+	}()
+	close(startCh)
+	if err := <-errCh; err != nil {
+		t.Fatalf("a worker got an error: %v", err)
+	}
+}
+
+func TestWARC_WriteSingleFile(t *testing.T) {
+	dir, err := ioutil.TempDir("", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(dir)
+
+	f, err := os.Create(filepath.Join(dir, "out.warc.gz"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	w := NewWriter(f)
+
+	writeManyRecords(t, w, 1000)
+	if err := w.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+}
+
+func TestWARC_WriteMulti(t *testing.T) {
+	dir, err := ioutil.TempDir("", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(dir)
+
+	var targetSize int64 = 10240
+	w, err := NewMultiWriter(filepath.Join(dir, "out.%s.warc.gz"), uint64(targetSize))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	writeManyRecords(t, w, 1000)
+	if err := w.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+
+	files, _ := ioutil.ReadDir(dir)
+	if len(files) < 2 {
+		t.Fatalf("MultiWriter didn't create enough files (%d)", len(files))
+	}
+	for _, f := range files[:len(files)-1] {
+		if f.Size() < targetSize {
+			t.Errorf("output file %s is too small (%d bytes)", f.Name(), f.Size())
+		}
+	}
+}
+
+func TestWARC_WriteMulti_Concurrent(t *testing.T) {
+	dir, err := ioutil.TempDir("", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(dir)
+
+	var targetSize int64 = 100000
+	w, err := NewMultiWriter(filepath.Join(dir, "out.%s.warc.gz"), uint64(targetSize))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	writeManyRecordsConcurrently(t, w, 1000, 10)
+	if err := w.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+
+	files, _ := ioutil.ReadDir(dir)
+	if len(files) < 2 {
+		t.Fatalf("MultiWriter didn't create enough files (%d)", len(files))
+	}
+	for _, f := range files[:len(files)-1] {
+		if f.Size() < targetSize {
+			t.Errorf("output file %s is too small (%d bytes)", f.Name(), f.Size())
+		}
+	}
+}