From 2f0eb59a4bdc0f18e9f7d2ef7fd7114c5007c887 Mon Sep 17 00:00:00 2001
From: Matthias Einwag <matthias.einwag@live.com>
Date: Tue, 11 Oct 2016 00:09:30 +0200
Subject: [PATCH] Fix the ping test by adding a delay to the write operation

---
 bench_test.go   |  3 +++
 session_test.go | 22 ++++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/bench_test.go b/bench_test.go
index 4377b83..b8fe46e 100644
--- a/bench_test.go
+++ b/bench_test.go
@@ -2,10 +2,13 @@ package yamux
 
 import (
 	"testing"
+	"time"
 )
 
 func BenchmarkPing(b *testing.B) {
 	client, server := testClientServer()
+	client.conn = &delayedConn{conn: client.conn, writeDelay: 10 * time.Millisecond}
+	server.conn = &delayedConn{conn: server.conn, writeDelay: 10 * time.Millisecond}
 	defer client.Close()
 	defer server.Close()
 
diff --git a/session_test.go b/session_test.go
index 0b4200e..076ef4f 100644
--- a/session_test.go
+++ b/session_test.go
@@ -51,6 +51,26 @@ func (p *pipeConn) Close() error {
 	return p.writer.Close()
 }
 
+type delayedConn struct {
+	conn       io.ReadWriteCloser
+	writeDelay time.Duration
+}
+
+func (c *delayedConn) Read(b []byte) (int, error) {
+	return c.conn.Read(b)
+}
+
+func (c *delayedConn) Write(b []byte) (int, error) {
+	if c.writeDelay != 0 {
+		time.Sleep(c.writeDelay)
+	}
+	return c.conn.Write(b)
+}
+
+func (c *delayedConn) Close() error {
+	return c.conn.Close()
+}
+
 func testConn() (io.ReadWriteCloser, io.ReadWriteCloser) {
 	read1, write1 := io.Pipe()
 	read2, write2 := io.Pipe()
@@ -86,6 +106,8 @@ func testClientServerConfig(conf *Config) (*Session, *Session) {
 
 func TestPing(t *testing.T) {
 	client, server := testClientServer()
+	client.conn = &delayedConn{conn: client.conn, writeDelay: 10 * time.Millisecond}
+	server.conn = &delayedConn{conn: server.conn, writeDelay: 10 * time.Millisecond}
 	defer client.Close()
 	defer server.Close()