From 2cbd0c4aefa628093584e9b608044c0f58e9373e Mon Sep 17 00:00:00 2001
From: Winni Neessen <wn@neessen.net>
Date: Mon, 23 Sep 2024 11:09:11 +0200
Subject: [PATCH] Add test cases for connection pool functionality

Added new test cases `TestConnPool_Get_Type` and `TestConnPool_Get` to verify connection pool operations. These tests ensure proper connection type and handling of pool size after connection retrieval and usage.
---
 connpool_test.go | 102 +++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 98 insertions(+), 4 deletions(-)

diff --git a/connpool_test.go b/connpool_test.go
index 589ad68..67a988f 100644
--- a/connpool_test.go
+++ b/connpool_test.go
@@ -8,6 +8,7 @@ import (
 	"context"
 	"fmt"
 	"net"
+	"sync"
 	"testing"
 	"time"
 )
@@ -30,6 +31,7 @@ func TestNewConnPool(t *testing.T) {
 	if err != nil {
 		t.Errorf("failed to create connection pool: %s", err)
 	}
+	defer pool.Close()
 	if pool == nil {
 		t.Errorf("connection pool is nil")
 		return
@@ -37,17 +39,109 @@ func TestNewConnPool(t *testing.T) {
 	if pool.Size() != 5 {
 		t.Errorf("expected 5 connections, got %d", pool.Size())
 	}
-	for i := 0; i < 5; i++ {
+	conn, err := pool.Get()
+	if err != nil {
+		t.Errorf("failed to get connection: %s", err)
+	}
+	if _, err := conn.Write([]byte("EHLO test.localhost.localdomain\r\nQUIT\r\n")); err != nil {
+		t.Errorf("failed to write quit command to first connection: %s", err)
+	}
+}
+
+func TestConnPool_Get_Type(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	serverPort := TestServerPortBase + 11
+	featureSet := "250-AUTH PLAIN\r\n250-8BITMIME\r\n250-DSN\r\n250 SMTPUTF8"
+	go func() {
+		if err := simpleSMTPServer(ctx, featureSet, true, serverPort); err != nil {
+			t.Errorf("failed to start test server: %s", err)
+			return
+		}
+	}()
+	time.Sleep(time.Millisecond * 300)
+
+	pool, err := newConnPool(serverPort)
+	if err != nil {
+		t.Errorf("failed to create connection pool: %s", err)
+	}
+	defer pool.Close()
+
+	conn, err := pool.Get()
+	if err != nil {
+		t.Errorf("failed to get new connection from pool: %s", err)
+		return
+	}
+
+	_, ok := conn.(*PoolConn)
+	if !ok {
+		t.Error("received connection from pool is not of type PoolConn")
+		return
+	}
+	if _, err := conn.Write([]byte("EHLO test.localhost.localdomain\r\nQUIT\r\n")); err != nil {
+		t.Errorf("failed to write quit command to first connection: %s", err)
+	}
+}
+
+func TestConnPool_Get(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	serverPort := TestServerPortBase + 12
+	featureSet := "250-AUTH PLAIN\r\n250-8BITMIME\r\n250-DSN\r\n250 SMTPUTF8"
+	go func() {
+		if err := simpleSMTPServer(ctx, featureSet, true, serverPort); err != nil {
+			t.Errorf("failed to start test server: %s", err)
+			return
+		}
+	}()
+	time.Sleep(time.Millisecond * 300)
+
+	p, _ := newConnPool(serverPort)
+	defer p.Close()
+
+	conn, err := p.Get()
+	if err != nil {
+		t.Errorf("failed to get new connection from pool: %s", err)
+		return
+	}
+	if _, err = conn.Write([]byte("EHLO test.localhost.localdomain\r\nQUIT\r\n")); err != nil {
+		t.Errorf("failed to write quit command to first connection: %s", err)
+	}
+
+	if p.Size() != 4 {
+		t.Errorf("getting new connection from pool failed. Expected pool size: 4, got %d", p.Size())
+	}
+
+	var wg sync.WaitGroup
+	for i := 0; i < 4; i++ {
+		wg.Add(1)
 		go func() {
-			conn, err := pool.Get()
+			defer wg.Done()
+			wgconn, err := p.Get()
 			if err != nil {
-				t.Errorf("failed to get connection: %s", err)
+				t.Errorf("failed to get new connection from pool: %s", err)
 			}
-			if _, err := conn.Write([]byte("EHLO test.localhost.localdomain\r\nQUIT\r\n")); err != nil {
+			if _, err = wgconn.Write([]byte("EHLO test.localhost.localdomain\r\nQUIT\r\n")); err != nil {
 				t.Errorf("failed to write quit command to first connection: %s", err)
 			}
 		}()
 	}
+	wg.Wait()
+
+	if p.Size() != 0 {
+		t.Errorf("Get error. Expecting 0, got %d", p.Size())
+	}
+
+	conn, err = p.Get()
+	if err != nil {
+		t.Errorf("failed to get new connection from pool: %s", err)
+	}
+	if _, err = conn.Write([]byte("EHLO test.localhost.localdomain\r\nQUIT\r\n")); err != nil {
+		t.Errorf("failed to write quit command to first connection: %s", err)
+	}
+	p.Close()
 }
 
 func newConnPool(port int) (Pool, error) {