Properly handle the case when servers read data by small chunks

1 files changed, 10 insertions, 6 deletions

diff --git a/server_timing_test.go b/server_timing_test.go
index 8d68890..8e237f3 100644
--- a/server_timing_test.go
+++ b/server_timing_test.go
@@ -154,25 +154,28 @@ type fakeServerConn struct {
 	net.TCPConn
 	ln            *fakeListener
 	requestsCount int
+	pos           int
 	closed        uint32
 }
 
 func (c *fakeServerConn) Read(b []byte) (int, error) {
 	nn := 0
-	for len(b) > len(c.ln.request) {
+	reqLen := len(c.ln.request)
+	for len(b) > 0 {
 		if c.requestsCount == 0 {
 			if nn == 0 {
 				return 0, io.EOF
 			}
 			return nn, nil
 		}
-		n := copy(b, c.ln.request)
+		pos := c.pos % reqLen
+		n := copy(b, c.ln.request[pos:])
 		b = b[n:]
 		nn += n
-		c.requestsCount--
-	}
-	if nn == 0 {
-		panic("server has too small buffer")
+		c.pos += n
+		if n+pos == reqLen {
+			c.requestsCount--
+		}
 	}
 	return nn, nil
 }
@@ -230,6 +233,7 @@ func (ln *fakeListener) Accept() (net.Conn, error) {
 	c := <-ln.ch
 	c.requestsCount = requestsCount
 	c.closed = 0
+	c.pos = 0
 
 	return c, nil
 }