test withconn

src/config_test.go

@@ -1,7 +1,13 @@
 package src
 
 import (
+	"fmt"
+	"io"
+	"log"
+	"net"
 	"os"
+	"strconv"
+	"sync"
 	"testing"
 )
 
@@ -30,3 +36,37 @@ func TestConfigLoadEnv(t *testing.T) {
 		t.Error("hello", c.Hello)
 	}
 }
+
+func TestConfigWithConn(t *testing.T) {
+	log.SetOutput(io.Discard)
+
+	c := Config{}
+
+	seen0, seen1 := false, false
+	conn0, conn1 := &net.IPConn{}, &net.IPConn{}
+	pool0, pool1 := &sync.Pool{}, &sync.Pool{}
+	pool0.Put(conn0)
+	pool1.Put(conn1)
+	c.forwards = []*sync.Pool{pool0, pool1}
+
+	for i := 0; i < 10; i++ {
+		i := strconv.Itoa(i)
+		t.Run(i, func(t *testing.T) {
+			if err := c.WithConn(i, func(c net.Conn) error {
+				c, _ = c.(*net.IPConn)
+				seen0 = seen0 || c == conn0
+				seen1 = seen1 || c == conn1
+				if c != conn0 && c != conn1 {
+					return fmt.Errorf("got unexepcted conn %v", c)
+				}
+				return nil
+			}); err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+
+	if !seen0 || !seen1 {
+		t.Fatalf("didnt see both conns")
+	}
+}