From b9d4b512c0154ee4d7f471ab106be1294b056957 Mon Sep 17 00:00:00 2001
From: Kyle Xiao <xiaoguoqiao@bytedance.com>
Date: Wed, 11 Sep 2024 13:37:16 +0800
Subject: [PATCH] test(thrift): simplified TestNocopyWrite

---
 internal/testutils/netpoll/netpoll.go |  1 +
 protocol/thrift/binary_test.go        | 36 ++++++++++-----------------
 2 files changed, 14 insertions(+), 23 deletions(-)

diff --git a/internal/testutils/netpoll/netpoll.go b/internal/testutils/netpoll/netpoll.go
index f09cb71..c5d4946 100644
--- a/internal/testutils/netpoll/netpoll.go
+++ b/internal/testutils/netpoll/netpoll.go
@@ -64,6 +64,7 @@ func (p *NetpollDirectWriter) WriteDirect(b []byte, remainCap int) error {
 	return nil
 }
 
+// WriteDirectN returns who many times WriteDirect called
 func (p *NetpollDirectWriter) WriteDirectN() int {
 	return len(p.wbuf)
 }
diff --git a/protocol/thrift/binary_test.go b/protocol/thrift/binary_test.go
index abe52e2..0544e4b 100644
--- a/protocol/thrift/binary_test.go
+++ b/protocol/thrift/binary_test.go
@@ -17,8 +17,6 @@
 package thrift
 
 import (
-	"bytes"
-	"math/rand"
 	"strings"
 	"testing"
 
@@ -422,32 +420,24 @@ func TestBinarySkip(t *testing.T) {
 }
 
 func TestNocopyWrite(t *testing.T) {
-	n := 100
-	// generate strings with the size above or below nocopyWriteThreshold
-	ss := make([]string, n)
-	for i := 0; i < len(ss); i++ {
-		c := string(rune('A' + i%61)) // 'A'(65) - '}"(125)
-		// 50% possibility will larger than nocopyWriteThreshold
-		ss[i] = strings.Repeat(c, 3*nocopyWriteThreshold/4+rand.Intn(nocopyWriteThreshold/2))
-	}
+	largestr := strings.Repeat("l", nocopyWriteThreshold)
+	smallstr := strings.Repeat("s", 10)
 
 	// generate expected data
 	x := BinaryProtocol{}
-	expectb := make([]byte, 0, n*2*nocopyWriteThreshold)
-	for _, s := range ss {
-		expectb = x.AppendString(expectb, s)
-	}
+	expectb := make([]byte, 0, 2*x.StringLength(smallstr)+x.StringLength(largestr))
+	expectb = x.AppendString(expectb, smallstr)
+	expectb = x.AppendString(expectb, largestr)
+	expectb = x.AppendString(expectb, smallstr)
 
 	// generate testing data
-	w := &netpoll.NetpollDirectWriter{}
 	i := 0
+	w := &netpoll.NetpollDirectWriter{}
 	b := w.Malloc(len(expectb))
-	for _, s := range ss {
-		i += x.WriteStringNocopy(b[i:], w, s)
-	}
-	b = w.Bytes()
-	wn := w.WriteDirectN()
-	require.True(t, wn > 0 && wn != n, wn)
-	require.Equal(t, len(expectb), len(b))
-	require.True(t, bytes.Equal(expectb, b))
+	i += x.WriteStringNocopy(b[i:], w, smallstr)
+	i += x.WriteStringNocopy(b[i:], w, largestr)
+	i += x.WriteStringNocopy(b[i:], w, smallstr)
+	require.Equal(t, len(expectb)-len(largestr), i) // without len(largestr)
+	require.Equal(t, 1, w.WriteDirectN())
+	require.Equal(t, expectb, w.Bytes())
 }