Make read benchmark smarter when extending file

1 files changed, 8 insertions, 3 deletions

diff --git a/main_test.go b/main_test.go
index 0dfa06a..667d626 100644
--- a/main_test.go
+++ b/main_test.go
@@ -136,14 +136,19 @@ func BenchmarkStreamRead(t *testing.B) {
 	buf := make([]byte, 1024*1024)
 	t.SetBytes(int64(len(buf)))
 
-	if t.N > 100 {
+	fn := plainDir + "BenchmarkWrite"
+	fi, _ := os.Stat(fn)
+	mb := int(fi.Size() / 1024 / 1024)
+
+	if t.N > mb {
 		// Grow file so we can satisfy the test
-		f2, err := os.OpenFile(plainDir + "BenchmarkWrite", os.O_WRONLY | os.O_APPEND, 0666)
+		fmt.Printf("Growing file to %d MB\n", t.N)
+		f2, err := os.OpenFile(fn, os.O_WRONLY | os.O_APPEND, 0666)
 		if err != nil {
 			fmt.Println(err)
 			t.FailNow()
 		}
-		for h := 0; h < t.N - 100 ; h++ {
+		for h := 0; h < t.N - mb ; h++ {
 			_, err = f2.Write(buf)
 			if err != nil {
 				fmt.Println(err)