1 files changed, 83 insertions, 0 deletions
diff --git a/reencode_test.go b/reencode_test.go
index 715e976..feabde5 100644
--- a/reencode_test.go
+++ b/reencode_test.go
@@ -5,6 +5,8 @@
 package lowmemjson
 
 import (
+	"errors"
+	"io"
 	"strings"
 	"testing"
 
@@ -240,3 +242,84 @@ func TestReEncoderStackSize(t *testing.T) {
 		assert.Equal(t, i+2, enc.stackSize())
 	}
 }
+
+var errNoSpace = errors.New("no space left on device")
+
+type limitedWriter struct {
+	Limit int
+	Inner io.Writer
+
+	n int
+}
+
+func (w *limitedWriter) Write(p []byte) (int, error) {
+	switch {
+	case w.n >= w.Limit:
+		return 0, errNoSpace
+	case w.n+len(p) > w.Limit:
+		n, err := w.Inner.Write(p[:w.Limit-w.n])
+		if n > 0 {
+			w.n += n
+		}
+		if err == nil {
+			err = errNoSpace
+		}
+		return n, err
+	default:
+		n, err := w.Inner.Write(p)
+		if n > 0 {
+			w.n += n
+		}
+		return n, err
+	}
+}
+
+func TestReEncodeIOErr(t *testing.T) {
+	t.Parallel()
+
+	input := `"😀"`
+	assert.Len(t, input, 6)
+
+	t.Run("bytes", func(t *testing.T) {
+		t.Parallel()
+
+		var out strings.Builder
+		enc := NewReEncoder(&limitedWriter{Limit: 5, Inner: &out}, ReEncoderConfig{})
+
+		n, err := enc.Write([]byte(input[:2]))
+		assert.NoError(t, err)
+		assert.Equal(t, 2, n)
+		// Of the 2 bytes "written", only one should be in
+		// `out` yet; the other should be in the UTF-8 buffer.
+		assert.Equal(t, input[:1], out.String())
+
+		n, err = enc.Write([]byte(input[2:]))
+		assert.ErrorIs(t, err, errNoSpace)
+		// Check that the byte in the UTF-8 buffer from the
+		// first .Write didn't count toward the total for this
+		// .Write.
+		assert.Equal(t, 3, n)
+		assert.Equal(t, input[:5], out.String())
+	})
+	t.Run("string", func(t *testing.T) {
+		t.Parallel()
+
+		var out strings.Builder
+		enc := NewReEncoder(&limitedWriter{Limit: 5, Inner: &out}, ReEncoderConfig{})
+
+		n, err := enc.WriteString(input[:2])
+		assert.NoError(t, err)
+		assert.Equal(t, 2, n)
+		// Of the 2 bytes "written", only one should be in
+		// `out` yet; the other should be in the UTF-8 buffer.
+		assert.Equal(t, input[:1], out.String())
+
+		n, err = enc.WriteString(input[2:])
+		assert.ErrorIs(t, err, errNoSpace)
+		// Check that the byte in the UTF-8 buffer from the
+		// first .Write didn't count toward the total for this
+		// .Write.
+		assert.Equal(t, 3, n)
+		assert.Equal(t, input[:5], out.String())
+	})
+}