From 1efac7fb67c8aba6c0d9d04d5ff4f5658cae29ca Mon Sep 17 00:00:00 2001
From: Winni Neessen <wn@neessen.net>
Date: Wed, 31 Jan 2024 11:22:01 +0100
Subject: [PATCH] Add new test cases for client functionalities and clean
 client code

Test cases have been added for numerous client functionalities including WithTLSPortPolicy option for the NewClient method, Client.SetSSLPort method, and the Client.DialWithContext method with the fallback port functionality. Minor code simplification has also been performed in the 'SetSSLPort' function in client.go file.
---
 client.go      |   8 ++--
 client_test.go | 104 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 107 insertions(+), 5 deletions(-)

diff --git a/client.go b/client.go
index a7c277d..2fbfe56 100644
--- a/client.go
+++ b/client.go
@@ -497,19 +497,17 @@ func (c *Client) SetSSL(s bool) {
 //
 // Port 465 is used when SSL set (true).
 // Port 25 is used when SSL is unset (false).
-// When the SSL connection fails and fallback is set to true,
+// When the SSL connection fails and fb is set to true,
 // the client will attempt to connect on port 25 using plaintext.
 func (c *Client) SetSSLPort(ssl bool, fb bool) {
+	c.port = DefaultPort
 	if ssl {
 		c.port = DefaultPortSSL
-	} else {
-		c.port = DefaultPort
 	}
 
+	c.fallbackPort = 0
 	if fb {
 		c.fallbackPort = DefaultPort
-	} else {
-		c.fallbackPort = 0
 	}
 
 	c.ssl = ssl
diff --git a/client_test.go b/client_test.go
index 5e39dc4..94381d9 100644
--- a/client_test.go
+++ b/client_test.go
@@ -238,6 +238,42 @@ func TestWithTLSPolicy(t *testing.T) {
 	}
 }
 
+// TestWithTLSPortPolicy tests the WithTLSPortPolicy() option for the NewClient() method
+func TestWithTLSPortPolicy(t *testing.T) {
+	tests := []struct {
+		name     string
+		value    TLSPolicy
+		want     string
+		wantPort int
+		fbPort   int
+		sf       bool
+	}{
+		{"Policy: TLSMandatory", TLSMandatory, TLSMandatory.String(), 587, 0, false},
+		{"Policy: TLSOpportunistic", TLSOpportunistic, TLSOpportunistic.String(), 587, 25, false},
+		{"Policy: NoTLS", NoTLS, NoTLS.String(), 25, 0, false},
+		{"Policy: Invalid", -1, "UnknownPolicy", 587, 0, true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c, err := NewClient(DefaultHost, WithTLSPortPolicy(tt.value))
+			if err != nil && !tt.sf {
+				t.Errorf("failed to create new client: %s", err)
+				return
+			}
+			if c.tlspolicy.String() != tt.want {
+				t.Errorf("failed to set TLSPortPolicy. Want: %s, got: %s", tt.want, c.tlspolicy)
+			}
+			if c.port != tt.wantPort {
+				t.Errorf("failed to set TLSPortPolicy, wanted port: %d, got: %d", tt.wantPort, c.port)
+			}
+			if c.fallbackPort != tt.fbPort {
+				t.Errorf("failed to set TLSPortPolicy, wanted fallbakc port: %d, got: %d", tt.fbPort,
+					c.fallbackPort)
+			}
+		})
+	}
+}
+
 // TestSetTLSPolicy tests the SetTLSPolicy() method for the Client object
 func TestSetTLSPolicy(t *testing.T) {
 	tests := []struct {
@@ -315,6 +351,42 @@ func TestSetSSL(t *testing.T) {
 	}
 }
 
+// TestSetSSLPort tests the Client.SetSSLPort method
+func TestClient_SetSSLPort(t *testing.T) {
+	tests := []struct {
+		name   string
+		value  bool
+		fb     bool
+		port   int
+		fbPort int
+	}{
+		{"SSL: on, fb: off", true, false, 465, 0},
+		{"SSL: on, fb: on", true, true, 465, 25},
+		{"SSL: off, fb: off", false, false, 25, 0},
+		{"SSL: off, fb: on", false, true, 25, 25},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c, err := NewClient(DefaultHost)
+			if err != nil {
+				t.Errorf("failed to create new client: %s", err)
+				return
+			}
+			c.SetSSLPort(tt.value, tt.fb)
+			if c.ssl != tt.value {
+				t.Errorf("failed to set SSL setting. Got: %t, want: %t", c.ssl, tt.value)
+			}
+			if c.port != tt.port {
+				t.Errorf("failed to set SSLPort, wanted port: %d, got: %d", c.port, tt.port)
+			}
+			if c.fallbackPort != tt.fbPort {
+				t.Errorf("failed to set SSLPort, wanted fallback port: %d, got: %d", c.fallbackPort,
+					tt.fbPort)
+			}
+		})
+	}
+}
+
 // TestSetUsername tests the SetUsername method for the Client object
 func TestSetUsername(t *testing.T) {
 	tests := []struct {
@@ -553,6 +625,38 @@ func TestClient_DialWithContext(t *testing.T) {
 	}
 }
 
+// TestClient_DialWithContext_Fallback tests the Client.DialWithContext method with the fallback
+// port functionality
+func TestClient_DialWithContext_Fallback(t *testing.T) {
+	c, err := getTestConnection(true)
+	if err != nil {
+		t.Skipf("failed to create test client: %s. Skipping tests", err)
+	}
+	c.SetTLSPortPolicy(TLSOpportunistic)
+	c.port = 999
+	ctx := context.Background()
+	if err := c.DialWithContext(ctx); err != nil {
+		t.Errorf("failed to dial with context: %s", err)
+		return
+	}
+	if c.co == nil {
+		t.Errorf("DialWithContext didn't fail but no connection found.")
+	}
+	if c.sc == nil {
+		t.Errorf("DialWithContext didn't fail but no SMTP client found.")
+	}
+	if err := c.Close(); err != nil {
+		t.Errorf("failed to close connection: %s", err)
+	}
+
+	c.port = 999
+	c.fallbackPort = 999
+	if err = c.DialWithContext(ctx); err == nil {
+		t.Error("dial with context was supposed to fail, but didn't")
+		return
+	}
+}
+
 // TestClient_DialWithContext_Debug tests the DialWithContext method for the Client object with debug
 // logging enabled on the SMTP client
 func TestClient_DialWithContext_Debug(t *testing.T) {