buffer/writer_test.go

package buffer_test

import (
	"code.squareroundforest.org/arpio/buffer"
	"errors"
	"io"
	"testing"
)

func TestWriter(t *testing.T) {
	t.Run("write out", func(t *testing.T) {
		w := &writer{}
		o := buffer.Options{BufferPool: buffer.NoPool(32)}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("456")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("789")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if err := b.Close(); err != nil {
			t.Fatal(err)
		}

		if string(w.written) != "123456789" {
			t.Fatal(string(w.written))
		}
	})

	t.Run("zero bytes when empty", func(t *testing.T) {
		w := &writer{}
		o := buffer.Options{BufferPool: buffer.NoPool(32)}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write(nil); n != 0 || err != nil {
			t.Fatal(n, err)
		}
	})

	t.Run("zero bytes when erred", func(t *testing.T) {
		w := &writer{errAfter: []int{3}}
		o := buffer.Options{BufferPool: buffer.NoPool(2)}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("456")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("789")); n != 0 || !errors.Is(err, errTest) {
			t.Fatal(n, err)
		}

		if n, err := b.Write(nil); n != 0 || !errors.Is(err, errTest) {
			t.Fatal(n, err)
		}

		if string(w.written) != "1234" {
			t.Fatal(string(w.written))
		}
	})

	t.Run("zero bytes when not empty", func(t *testing.T) {
		w := &writer{errAfter: []int{3}}
		o := buffer.Options{BufferPool: buffer.NoPool(32)}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write(nil); n != 0 || err != nil {
			t.Fatal(n, err)
		}

		if len(w.written) != 0 {
			t.Fatal(string(w.written))
		}
	})

	t.Run("get buffer fails", func(t *testing.T) {
		w := &writer{}
		p := &fakePool{
			allocSize: 32,
			errAfter:  []int{0},
		}

		o := buffer.Options{BufferPool: p}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); n != 0 || !errors.Is(err, errTest) {
			t.Fatal(n, err)
		}
	})

	t.Run("no underlying write until full", func(t *testing.T) {
		w := &writer{}
		o := buffer.Options{BufferPool: buffer.NoPool(4)}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if len(w.written) != 0 {
			t.Fatal(string(w.written))
		}

		if n, err := b.Write([]byte("456")); n != 3 || err != nil {
			t.Fatal(n, err)
		}
		if string(w.written) != "1234" {
			t.Fatal(string(w.written))
		}

		if n, err := b.Write([]byte("789")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if string(w.written) != "12345678" {
			t.Fatal(string(w.written))
		}

		if err := b.Close(); err != nil {
			t.Fatal(err)
		}

		if string(w.written) != "123456789" {
			t.Fatal(string(w.written))
		}
	})

	t.Run("auto flush when write larger than buffer", func(t *testing.T) {
		w := &writer{}
		o := buffer.Options{BufferPool: buffer.NoPool(4)}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123456")); err != nil {
			t.Fatal(n, err)
		}

		if string(w.written) != "1234" {
			t.Fatal(string(w.written))
		}
	})

	t.Run("auto flush when multiple smaller writes fill the buffer", func(t *testing.T) {
		w := &writer{}
		o := buffer.Options{BufferPool: buffer.NoPool(4)}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); err != nil {
			t.Fatal(n, err)
		}

		if len(w.written) != 0 {
			t.Fatal(string(w.written))
		}

		if n, err := b.Write([]byte("456")); err != nil {
			t.Fatal(n, err)
		}

		if string(w.written) != "1234" {
			t.Fatal(string(w.written))
		}
	})

	t.Run("flush on close", func(t *testing.T) {
		w := &writer{}
		o := buffer.Options{BufferPool: buffer.NoPool(4)}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); err != nil {
			t.Fatal(n, err)
		}

		if len(w.written) != 0 {
			t.Fatal(string(w.written))
		}

		if err := b.Close(); err != nil {
			t.Fatal(err)
		}

		if string(w.written) != "123" {
			t.Fatal(string(w.written))
		}
	})

	t.Run("manual flush", func(t *testing.T) {
		w := &writer{}
		o := buffer.Options{BufferPool: buffer.NoPool(4)}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); err != nil {
			t.Fatal(n, err)
		}

		if len(w.written) != 0 {
			t.Fatal(string(w.written))
		}

		if err := b.Flush(); err != nil {
			t.Fatal(err)
		}

		if string(w.written) != "123" {
			t.Fatal(string(w.written))
		}
	})

	t.Run("write after flush", func(t *testing.T) {
		w := &writer{}
		o := buffer.Options{BufferPool: buffer.NoPool(4)}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); err != nil {
			t.Fatal(n, err)
		}

		if len(w.written) != 0 {
			t.Fatal(string(w.written))
		}

		if err := b.Flush(); err != nil {
			t.Fatal(err)
		}

		if string(w.written) != "123" {
			t.Fatal(string(w.written))
		}

		if n, err := b.Write([]byte("456")); err != nil {
			t.Fatal(n, err)
		}

		if err := b.Close(); err != nil {
			t.Fatal(err)
		}

		if string(w.written) != "123456" {
			t.Fatal(string(w.written))
		}
	})

	t.Run("zero write recover", func(t *testing.T) {
		w := &writer{zeroAfter: []int{2}}
		p := &fakePool{allocSize: 2}
		o := buffer.Options{BufferPool: p}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("456")); err != nil {
			t.Fatal(n, err)
		}

		if err := b.Close(); err != nil {
			t.Fatal(err)
		}

		if string(w.written) != "123456" {
			t.Fatal(string(w.written))
		}
	})

	t.Run("zero write terminal", func(t *testing.T) {
		w := &writer{zeroAfter: []int{2, 2}}
		p := &fakePool{allocSize: 2}
		o := buffer.Options{BufferPool: p}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("456")); n != 1 || !errors.Is(err, io.ErrShortWrite) {
			t.Fatal(n, err)
		}
	})

	t.Run("partial write", func(t *testing.T) {
		w := &writer{shortAfter: []int{2}}
		p := &fakePool{allocSize: 2}
		o := buffer.Options{BufferPool: p}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("456")); err != nil {
			t.Fatal(n, err)
		}

		if err := b.Close(); err != nil {
			t.Fatal(err)
		}

		if string(w.written) != "123456" {
			t.Fatal(string(w.written))
		}
	})

	t.Run("buffer released on write error immediately", func(t *testing.T) {
		w := &writer{errAfter: []int{0}}
		p := &fakePool{allocSize: 2}
		o := buffer.Options{BufferPool: p}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); n != 2 || !errors.Is(err, errTest) {
			t.Fatal(n, err)
		}

		if p.alloc != 1 && p.free != 1 {
			t.Fatal(p.alloc, p.free)
		}
	})

	t.Run("buffer released on write error", func(t *testing.T) {
		w := &writer{errAfter: []int{3}}
		p := &fakePool{allocSize: 2}
		o := buffer.Options{BufferPool: p}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("456")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("789")); n != 0 || !errors.Is(err, errTest) {
			t.Fatal(n, err)
		}

		if p.alloc != 1 && p.free != 1 {
			t.Fatal(p.alloc, p.free)
		}
	})

	t.Run("close on err", func(t *testing.T) {
		w := &writer{errAfter: []int{3}}
		p := &fakePool{allocSize: 2}
		o := buffer.Options{BufferPool: p}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("456")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("789")); n != 0 || !errors.Is(err, errTest) {
			t.Fatal(n, err)
		}

		if p.alloc != 1 && p.free != 1 {
			t.Fatal(p.alloc, p.free)
		}

		if err := b.Close(); err != nil {
			t.Fatal(err)
		}
	})

	t.Run("close and flush err", func(t *testing.T) {
		w := &writer{errAfter: []int{3}}
		p := &fakePool{allocSize: 2}
		o := buffer.Options{BufferPool: p}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("456")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if err := b.Close(); !errors.Is(err, errTest) {
			t.Fatal(err)
		}
	})

	t.Run("close", func(t *testing.T) {
		w := &writer{}
		p := &fakePool{allocSize: 2}
		o := buffer.Options{BufferPool: p}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if n, err := b.Write([]byte("456")); n != 3 || err != nil {
			t.Fatal(n, err)
		}

		if err := b.Close(); err != nil {
			t.Fatal(err)
		}

		if string(w.written) != "123456" {
			t.Fatal(string(w.written))
		}
	})

	t.Run("zero allocation", func(t *testing.T) {
		w := &writer{}
		p := &fakePool{}
		o := buffer.Options{BufferPool: p}
		b := buffer.BufferedWriter(w, o)
		if n, err := b.Write([]byte("123")); n != 0 || !errors.Is(err, buffer.ErrZeroAllocation) {
			t.Fatal(n, err)
		}
	})

	t.Run("read from", func(t *testing.T) {
		t.Run("read out", func(t *testing.T) {
			w := &writer{}
			r := &gen{max: 1 << 12}
			p := buffer.NoPool(1 << 9)
			o := buffer.Options{BufferPool: p}
			b := buffer.BufferedWriter(w, o)
			n, err := b.ReadFrom(r)
			if n != 1<<12 || err != nil {
				t.Fatal(n, err)
			}
		})

		t.Run("read out final flush fail", func(t *testing.T) {
			w := &writer{errAfter: []int{1<<12 - 1<<9}}
			r := &gen{
				max:     1 << 12,
				fastErr: true,
			}

			p := buffer.NoPool(1 << 9)
			o := buffer.Options{BufferPool: p}
			b := buffer.BufferedWriter(w, o)
			n, err := b.ReadFrom(r)
			if n != 1<<12 || !errors.Is(err, errTest) {
				t.Fatal(n, err)
			}
		})

		t.Run("read error", func(t *testing.T) {
			w := &writer{}
			r := &gen{
				max:      1 << 12,
				errAfter: []int{1 << 11},
			}

			p := buffer.NoPool(1 << 9)
			o := buffer.Options{BufferPool: p}
			b := buffer.BufferedWriter(w, o)
			n, err := b.ReadFrom(r)
			if n != 1<<11 || !errors.Is(err, errTest) {
				t.Fatal(n, err)
			}
		})

		t.Run("write error", func(t *testing.T) {
			w := &writer{errAfter: []int{1 << 11}}
			r := &gen{max: 1 << 12}
			p := buffer.NoPool(1 << 9)
			o := buffer.Options{BufferPool: p}
			b := buffer.BufferedWriter(w, o)
			n, err := b.ReadFrom(r)
			if n != 1<<11+1<<9 || !errors.Is(err, errTest) {
				t.Fatal(n, err)
			}
		})

		t.Run("zero allocation", func(t *testing.T) {
			w := &writer{}
			r := &gen{max: 1 << 12}
			p := &fakePool{}
			o := buffer.Options{BufferPool: p}
			b := buffer.BufferedWriter(w, o)
			n, err := b.ReadFrom(r)
			if n != 0 || !errors.Is(err, buffer.ErrZeroAllocation) {
				t.Fatal(n, err)
			}
		})

		t.Run("no progress", func(t *testing.T) {
			w := &writer{}
			r := &gen{
				max:           1 << 15,
				nullReadAfter: []int{256, 256},
			}

			p := &fakePool{allocSize: 1 << 9}
			o := buffer.Options{BufferPool: p}
			b := buffer.BufferedWriter(w, o)
			n, err := b.ReadFrom(r)
			if n != 512 || !errors.Is(err, io.ErrNoProgress) {
				t.Fatal(n, err)
			}
		})
	})
}