From c55c4aa365bce9de1fe3c1659e98a25864441fb0 Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Tue, 20 Aug 2024 09:46:04 +0200
Subject: [PATCH] fix for non-deterministic test.

---
 pkg/testsupport/channelreadwritecloser.go  | 10 ++++++++++
 pkg/testsupport/inmemoryconnection.go      |  3 ++-
 pkg/testsupport/inmemoryconnection_test.go |  1 +
 3 files changed, 13 insertions(+), 1 deletion(-)
 create mode 100644 pkg/testsupport/inmemoryconnection_test.go

diff --git a/pkg/testsupport/channelreadwritecloser.go b/pkg/testsupport/channelreadwritecloser.go
index c10e639..c5e9d43 100644
--- a/pkg/testsupport/channelreadwritecloser.go
+++ b/pkg/testsupport/channelreadwritecloser.go
@@ -51,7 +51,17 @@ func (rw *ChannelReadWriter) Write(p []byte) (n int, err error) {
 	if rw.closed {
 		return 0, errors.New("Write on closed channel")
 	}
+
+	// if context is canceled it should never write
 	select {
+	case <-rw.ctx.Done():
+		rw.Close()
+		return 0, io.ErrClosedPipe
+	default:
+	}
+
+	select {
+	// deal with closing duirng the write
 	case <-rw.ctx.Done():
 		rw.Close()
 		return 0, io.ErrClosedPipe
diff --git a/pkg/testsupport/inmemoryconnection.go b/pkg/testsupport/inmemoryconnection.go
index 502d819..2739f0e 100644
--- a/pkg/testsupport/inmemoryconnection.go
+++ b/pkg/testsupport/inmemoryconnection.go
@@ -11,7 +11,8 @@ type InmemoryConnection struct {
 
 func NewInmemoryConnection(ctx context.Context, addr string) *InmemoryConnection {
 	pipe := InmemoryConnection{
-		ctx:         ctx,
+		ctx: ctx,
+		// arbitrary unbuffered channel, unbuffered is more similar to TCP connections.
 		frontToBack: make(chan []byte),
 		backToFront: make(chan []byte),
 		addr:        addr,
diff --git a/pkg/testsupport/inmemoryconnection_test.go b/pkg/testsupport/inmemoryconnection_test.go
new file mode 100644
index 0000000..186d85f
--- /dev/null
+++ b/pkg/testsupport/inmemoryconnection_test.go
@@ -0,0 +1 @@
+package testsupport