buffer/io_test.go

package buffer_test

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

type createReader func(r io.Reader, o buffer.Options) buffer.Reader

type gen struct {
	rng                []byte
	max                int
	fastErr            bool
	nullReadAfter      []int
	errAfter           []int
	customContentAfter []int
	customContent      map[int][]byte
	counter            int
}

type writer struct {
	written    []byte
	errAfter   []int
	shortAfter []int
}

var (
	genRange    = []byte("abcdefghi")
	utf8Range   = []byte("aábéícóöődúüeű")
	utf8W2Range = []byte("áéíóöőúüű")
	errTest     = errors.New("test error")
	errTest2    = errors.New("test error 2")
)

func (g *gen) Read(p []byte) (int, error) {
	if g.max == 0 {
		return 0, io.EOF
	}

	if len(g.nullReadAfter) > 0 && g.counter >= g.nullReadAfter[0] {
		g.nullReadAfter = g.nullReadAfter[1:]
		return 0, nil
	}

	if len(g.errAfter) > 0 && g.counter >= g.errAfter[0] {
		g.errAfter = g.errAfter[1:]
		return 0, errTest
	}

	l := len(p)
	hasMax := g.max > 0
	if hasMax && l > g.max {
		l = g.max
	}

	if len(g.rng) == 0 {
		g.rng = genRange
	}

	var n int
	for l > 0 {
		rng := make([]byte, len(g.rng))
		copy(rng, g.rng)
		c := g.counter % len(rng)
		rng = append(rng[c:], rng[:c]...)
		li := l
		if li > len(rng) {
			li = len(rng)
		}

		cc := len(g.customContentAfter) > 0 &&
			g.counter <= g.customContentAfter[0] &&
			g.counter+li > g.customContentAfter[0]
		if cc && g.counter != g.customContentAfter[0] {
			li = g.customContentAfter[0] - g.counter
			cc = false
		}

		var ni int
		if cc {
			content := g.customContent[g.customContentAfter[0]]
			ni = copy(p[:li], content)
			if ni == len(content) {
				g.customContentAfter = g.customContentAfter[1:]
			} else {
				g.customContentAfter[0] += ni
				g.customContent[g.customContentAfter[0]] = content[ni:]
			}
		} else {
			ni = copy(p[:li], rng)
		}

		n += ni
		p = p[ni:]
		l -= ni
		g.counter += ni
		if cc {
			c := g.counter % len(g.rng)
			rng = append(rng[c:], rng[:c]...)
		}
	}

	if hasMax {
		g.max -= n
	}

	if len(g.errAfter) > 0 && g.counter >= g.errAfter[0] && g.fastErr {
		g.errAfter = g.errAfter[1:]
		return n, errTest
	}

	if hasMax && g.max == 0 && g.fastErr {
		return n, io.EOF
	}

	return n, nil
}

func generateFrom(rng []byte, n int) []byte {
	g := &gen{
		rng: rng,
		max: n,
	}

	b, _ := io.ReadAll(g)
	return b
}

func generate(n int) []byte {
	return generateFrom(genRange, n)
}

func (w *writer) Write(p []byte) (int, error) {
	if len(w.errAfter) > 0 && len(w.written) >= w.errAfter[0] {
		w.errAfter = w.errAfter[1:]
		return 0, errTest
	}

	if len(p) > 0 && len(w.shortAfter) > 0 && len(w.written) >= w.shortAfter[0] {
		w.shortAfter = w.shortAfter[1:]
		p = p[:len(p)/2]
	}

	wp := make([]byte, len(p))
	copy(wp, p)
	w.written = append(w.written, wp...)
	return len(p), nil
}

func testContent(r io.Reader, o buffer.Options) buffer.Reader {
	return buffer.BufferedContent(
		buffer.ContentFunc(func(w io.Writer) (int64, error) {
			return io.Copy(w, r)
		}),
		o,
	)
}