From 0dcede645782184eaff390285d7204b37b0f0ed9 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Sat, 16 Sep 2023 22:33:30 +0800
Subject: [PATCH] chore: refactor log limit in rest (#3572)

---
 core/iox/nopcloser_test.go | 12 ++++++++++++
 core/iox/read.go           | 14 +++++++-------
 core/iox/read_test.go      | 10 ++++++++++
 core/iox/tee.go            |  2 ++
 core/iox/tee_test.go       | 19 ++++++++++++-------
 core/iox/textfile.go       |  3 ++-
 core/iox/textfile_test.go  |  5 +++++
 7 files changed, 50 insertions(+), 15 deletions(-)
 create mode 100644 core/iox/nopcloser_test.go

diff --git a/core/iox/nopcloser_test.go b/core/iox/nopcloser_test.go
new file mode 100644
index 00000000..5dfd4d56
--- /dev/null
+++ b/core/iox/nopcloser_test.go
@@ -0,0 +1,12 @@
+package iox
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNopCloser(t *testing.T) {
+	closer := NopCloser(nil)
+	assert.NoError(t, closer.Close())
+}
diff --git a/core/iox/read.go b/core/iox/read.go
index 1f27f445..3ec024ec 100644
--- a/core/iox/read.go
+++ b/core/iox/read.go
@@ -28,6 +28,13 @@ func DupReadCloser(reader io.ReadCloser) (io.ReadCloser, io.ReadCloser) {
 	return io.NopCloser(tee), io.NopCloser(&buf)
 }
 
+// KeepSpace customizes the reading functions to keep leading and tailing spaces.
+func KeepSpace() TextReadOption {
+	return func(o *textReadOptions) {
+		o.keepSpace = true
+	}
+}
+
 // LimitDupReadCloser returns two io.ReadCloser that read from the first will be written to the second.
 // But the second io.ReadCloser is limited to up to n bytes.
 // The first returned reader needs to be read first, because the content
@@ -38,13 +45,6 @@ func LimitDupReadCloser(reader io.ReadCloser, n int64) (io.ReadCloser, io.ReadCl
 	return io.NopCloser(tee), io.NopCloser(&buf)
 }
 
-// KeepSpace customizes the reading functions to keep leading and tailing spaces.
-func KeepSpace() TextReadOption {
-	return func(o *textReadOptions) {
-		o.keepSpace = true
-	}
-}
-
 // ReadBytes reads exactly the bytes with the length of len(buf)
 func ReadBytes(reader io.Reader, buf []byte) error {
 	var got int
diff --git a/core/iox/read_test.go b/core/iox/read_test.go
index b31ce283..d2eb2435 100644
--- a/core/iox/read_test.go
+++ b/core/iox/read_test.go
@@ -51,6 +51,11 @@ b`,
 	}
 }
 
+func TestReadTextError(t *testing.T) {
+	_, err := ReadText("not-exist")
+	assert.NotNil(t, err)
+}
+
 func TestReadTextLines(t *testing.T) {
 	text := `1
 
@@ -94,6 +99,11 @@ func TestReadTextLines(t *testing.T) {
 	}
 }
 
+func TestReadTextLinesError(t *testing.T) {
+	_, err := ReadTextLines("not-exist")
+	assert.NotNil(t, err)
+}
+
 func TestDupReadCloser(t *testing.T) {
 	input := "hello"
 	reader := io.NopCloser(bytes.NewBufferString(input))
diff --git a/core/iox/tee.go b/core/iox/tee.go
index 07b00158..11c5221e 100644
--- a/core/iox/tee.go
+++ b/core/iox/tee.go
@@ -27,7 +27,9 @@ func (t *limitTeeReader) Read(p []byte) (n int, err error) {
 		if n, err := t.w.Write(p[:limit]); err != nil {
 			return n, err
 		}
+
 		t.n -= limit
 	}
+
 	return
 }
diff --git a/core/iox/tee_test.go b/core/iox/tee_test.go
index 56e457c7..87f5267e 100644
--- a/core/iox/tee_test.go
+++ b/core/iox/tee_test.go
@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"io"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestLimitTeeReader(t *testing.T) {
@@ -22,14 +24,17 @@ func TestLimitTeeReader(t *testing.T) {
 	if !bytes.Equal(wb.Bytes(), src[:limit]) {
 		t.Errorf("bytes written = %q want %q", wb.Bytes(), src)
 	}
-	if n, err := r.Read(dst); n != 0 || err != io.EOF {
-		t.Errorf("r.Read at EOF = %d, %v want 0, EOF", n, err)
-	}
+
+	n, err := r.Read(dst)
+	assert.Equal(t, 0, n)
+	assert.Equal(t, io.EOF, err)
+
 	rb = bytes.NewBuffer(src)
 	pr, pw := io.Pipe()
-	pr.Close()
-	r = LimitTeeReader(rb, pw, limit)
-	if n, err := io.ReadFull(r, dst); n != 0 || err != io.ErrClosedPipe {
-		t.Errorf("closed tee: ReadFull(r, dst) = %d, %v; want 0, EPIPE", n, err)
+	if assert.NoError(t, pr.Close()) {
+		r = LimitTeeReader(rb, pw, limit)
+		n, err := io.ReadFull(r, dst)
+		assert.Equal(t, 0, n)
+		assert.Equal(t, io.ErrClosedPipe, err)
 	}
 }
diff --git a/core/iox/textfile.go b/core/iox/textfile.go
index 8feced43..89083cea 100644
--- a/core/iox/textfile.go
+++ b/core/iox/textfile.go
@@ -2,6 +2,7 @@ package iox
 
 import (
 	"bytes"
+	"errors"
 	"io"
 	"os"
 )
@@ -26,7 +27,7 @@ func CountLines(file string) (int, error) {
 		count += bytes.Count(buf[:c], lineSep)
 
 		switch {
-		case err == io.EOF:
+		case errors.Is(err, io.EOF):
 			if noEol {
 				count++
 			}
diff --git a/core/iox/textfile_test.go b/core/iox/textfile_test.go
index 011d674b..c58b3048 100644
--- a/core/iox/textfile_test.go
+++ b/core/iox/textfile_test.go
@@ -24,3 +24,8 @@ func TestCountLines(t *testing.T) {
 	assert.Nil(t, err)
 	assert.Equal(t, 4, lines)
 }
+
+func TestCountLinesError(t *testing.T) {
+	_, err := CountLines("not-exist")
+	assert.NotNil(t, err)
+}