diff options
author | Luke Shumaker <lukeshu@lukeshu.com> | 2023-02-25 16:17:06 -0700 |
---|---|---|
committer | Luke Shumaker <lukeshu@lukeshu.com> | 2023-02-25 18:42:05 -0700 |
commit | 7bd0072b5896bfc4172b6bda778cf149dd6282fa (patch) | |
tree | 518c9e62aace98fcca449fe73996abb8ee769db6 /reencode_test.go | |
parent | 4233e5012ece6d5a7fee3b5a518c41d916e1cf52 (diff) |
reencode: Fix the byte count for partial writes
Diffstat (limited to 'reencode_test.go')
-rw-r--r-- | reencode_test.go | 83 |
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()) + }) +} |