Repository: inbucket/inbucket
Pull request #289 (2022-08-25 04:21:27)

Refactor SMTP handler_test log collection
Debugging SMTP handler sessions is difficult, as all the logs for all tests were being sent to the terminal when any test failed.  This change uses zerolog functionality to collect logs for a particular test and log it to testing.T.

- smtp: allow logger to be passed into startSession
- smtp: remove logbuf and teardown from handler_test
- smtp: handler_test log output to testing.T



Patch:
diff --git a/pkg/server/smtp/handler.go b/pkg/server/smtp/handler.go
index f9e9a55..cc1c20a 100644
--- a/pkg/server/smtp/handler.go
+++ b/pkg/server/smtp/handler.go
@@ -15,7 +15,6 @@ import (
 
 	"github.com/inbucket/inbucket/pkg/policy"
 	"github.com/rs/zerolog"
-	"github.com/rs/zerolog/log"
 )
 
 // State tracks the current mode of our SMTP state machine.
@@ -144,8 +143,8 @@ func (s *Session) String() string {
  *  4. If bad cmd, respond error
  *  5. Goto 2
  */
-func (s *Server) startSession(id int, conn net.Conn) {
-	logger := log.Hook(logHook{}).With().
+func (s *Server) startSession(id int, conn net.Conn, logger zerolog.Logger) {
+	logger = logger.Hook(logHook{}).With().
 		Str("module", "smtp").
 		Str("remote", conn.RemoteAddr().String()).
 		Int("session", id).Logger()
diff --git a/pkg/server/smtp/listener.go b/pkg/server/smtp/listener.go
index 5a281f7..0a38282 100644
--- a/pkg/server/smtp/listener.go
+++ b/pkg/server/smtp/listener.go
@@ -170,7 +170,7 @@ func (s *Server) serve(ctx context.Context) {
 			tempDelay = 0
 			expConnectsTotal.Add(1)
 			s.wg.Add(1)
-			go s.startSession(sessionID, conn)
+			go s.startSession(sessionID, conn, log.Logger)
 		}
 	}
 }


Test patch:
diff --git a/pkg/server/smtp/handler_test.go b/pkg/server/smtp/handler_test.go
index f07e824..e1f67e1 100644
--- a/pkg/server/smtp/handler_test.go
+++ b/pkg/server/smtp/handler_test.go
@@ -1,14 +1,11 @@
 package smtp
 
 import (
-	"bytes"
 	"fmt"
 	"io"
 
-	"log"
 	"net"
 	"net/textproto"
-	"os"
 	"testing"
 	"time"
 
@@ -17,6 +14,7 @@ import (
 	"github.com/inbucket/inbucket/pkg/policy"
 	"github.com/inbucket/inbucket/pkg/storage"
 	"github.com/inbucket/inbucket/pkg/test"
+	"github.com/rs/zerolog"
 )
 
 type scriptStep struct {
@@ -27,8 +25,8 @@ type scriptStep struct {
 // Test commands in GREET state
 func TestGreetState(t *testing.T) {
 	ds := test.NewStore()
-	server, logbuf, teardown := setupSMTPServer(ds)
-	defer teardown()
+	server := setupSMTPServer(ds)
+	defer server.Drain() // Required to prevent test logging data race.
 
 	// Test out some mangled HELOs
 	script := []scriptStep{
@@ -77,19 +75,13 @@ func TestGreetState(t *testing.T) {
 		t.Error(err)
 	}
 
-	if t.Failed() {
-		// Wait for handler to finish logging
-		time.Sleep(2 * time.Second)
-		// Dump buffered log data if there was a failure
-		_, _ = io.Copy(os.Stderr, logbuf)
-	}
 }
 
 // Test commands in READY state
 func TestEmptyEnvelope(t *testing.T) {
 	ds := test.NewStore()
-	server, logbuf, teardown := setupSMTPServer(ds)
-	defer teardown()
+	server := setupSMTPServer(ds)
+	defer server.Drain()
 
 	// Test out some empty envelope without blanks
 	script := []scriptStep{
@@ -97,8 +89,6 @@ func TestEmptyEnvelope(t *testing.T) {
 		{"MAIL FROM:<>", 250},
 	}
 	if err := playSession(t, server, script); err != nil {
-		// Dump buffered log data if there was a failure
-		_, _ = io.Copy(os.Stderr, logbuf)
 		t.Error(err)
 	}
 
@@ -108,8 +98,6 @@ func TestEmptyEnvelope(t *testing.T) {
 		{"MAIL FROM: <>", 250},
 	}
 	if err := playSession(t, server, script); err != nil {
-		// Dump buffered log data if there was a failure
-		_, _ = io.Copy(os.Stderr, logbuf)
 		t.Error(err)
 	}
 }
@@ -117,8 +105,8 @@ func TestEmptyEnvelope(t *testing.T) {
 // Test AUTH
 func TestAuth(t *testing.T) {
 	ds := test.NewStore()
-	server, logbuf, teardown := setupSMTPServer(ds)
-	defer teardown()
+	server := setupSMTPServer(ds)
+	defer server.Drain()
 
 	// PLAIN AUTH
 	script := []scriptStep{
@@ -149,20 +137,13 @@ func TestAuth(t *testing.T) {
 	if err := playSession(t, server, script); err != nil {
 		t.Error(err)
 	}
-
-	if t.Failed() {
-		// Wait for handler to finish logging
-		time.Sleep(2 * time.Second)
-		// Dump buffered log data if there was a failure
-		_, _ = io.Copy(os.Stderr, logbuf)
-	}
 }
 
 // Test commands in READY state
 func TestReadyState(t *testing.T) {
 	ds := test.NewStore()
-	server, logbuf, teardown := setupSMTPServer(ds)
-	defer teardown()
+	server := setupSMTPServer(ds)
+	defer server.Drain()
 
 	// Test out some mangled READY commands
 	script := []scriptStep{
@@ -221,20 +202,13 @@ func TestReadyState(t *testing.T) {
 	if err := playSession(t, server, script); err != nil {
 		t.Error(err)
 	}
-
-	if t.Failed() {
-		// Wait for handler to finish logging
-		time.Sleep(2 * time.Second)
-		// Dump buffered log data if there was a failure
-		_, _ = io.Copy(os.Stderr, logbuf)
-	}
 }
 
 // Test commands in MAIL state
 func TestMailState(t *testing.T) {
 	mds := test.NewStore()
-	server, logbuf, teardown := setupSMTPServer(mds)
-	defer teardown()
+	server := setupSMTPServer(mds)
+	defer server.Drain()
 
 	// Test out some mangled READY commands
 	script := []scriptStep{
@@ -335,23 +309,16 @@ func TestMailState(t *testing.T) {
 	if err := playSession(t, server, script); err != nil {
 		t.Error(err)
 	}
-
-	if t.Failed() {
-		// Wait for handler to finish logging
-		time.Sleep(2 * time.Second)
-		// Dump buffered log data if there was a failure
-		_, _ = io.Copy(os.Stderr, logbuf)
-	}
 }
 
 // Test commands in DATA state
 func TestDataState(t *testing.T) {
 	mds := test.NewStore()
-	server, logbuf, teardown := setupSMTPServer(mds)
-	defer teardown()
+	server := setupSMTPServer(mds)
+	defer server.Drain()
 
 	var script []scriptStep
-	pipe := setupSMTPSession(server)
+	pipe := setupSMTPSession(t, server)
 	c := textproto.NewConn(pipe)
 
 	if code, _, err := c.ReadCodeLine(220); err != nil {
@@ -366,6 +333,7 @@ func TestDataState(t *testing.T) {
 	if err := playScriptAgainst(t, c, script); err != nil {
 		t.Error(err)
 	}
+
 	// Send a message
 	body := `To: u1@gmail.com
 From: john@gmail.com
@@ -379,9 +347,11 @@ Hi!
 	if code, _, err := c.ReadCodeLine(250); err != nil {
 		t.Errorf("Expected a 250 greeting, got %v", code)
 	}
+	_, _ = c.Cmd("QUIT")
+	_, _, _ = c.ReadCodeLine(221)
 
 	// Test with no useful headers.
-	pipe = setupSMTPSession(server)
+	pipe = setupSMTPSession(t, server)
 	c = textproto.NewConn(pipe)
 	if code, _, err := c.ReadCodeLine(220); err != nil {
 		t.Errorf("Expected a 220 greeting, got %v", code)
@@ -395,29 +365,25 @@ Hi!
 	if err := playScriptAgainst(t, c, script); err != nil {
 		t.Error(err)
 	}
+
 	// Send a message
 	body = `X-Useless-Header: true
 
-Hi! Can you still deliver this?
-`
+	Hi! Can you still deliver this?
+	`
 	dw = c.DotWriter()
 	_, _ = io.WriteString(dw, body)
 	_ = dw.Close()
 	if code, _, err := c.ReadCodeLine(250); err != nil {
 		t.Errorf("Expected a 250 greeting, got %v", code)
 	}
-
-	if t.Failed() {
-		// Wait for handler to finish logging
-		time.Sleep(2 * time.Second)
-		// Dump buffered log data if there was a failure
-		_, _ = io.Copy(os.Stderr, logbuf)
-	}
+	_, _ = c.Cmd("QUIT")
+	_, _, _ = c.ReadCodeLine(221)
 }
 
 // playSession creates a new session, reads the greeting and then plays the script
 func playSession(t *testing.T, server *Server, script []scriptStep) error {
-	pipe := setupSMTPSession(server)
+	pipe := setupSMTPSession(t, server)
 	c := textproto.NewConn(pipe)
 
 	if code, _, err := c.ReadCodeLine(220); err != nil {
@@ -467,7 +433,7 @@ func (m *mockConn) SetDeadline(t time.Time) error      { return nil }
 func (m *mockConn) SetReadDeadline(t time.Time) error  { return nil }
 func (m *mockConn) SetWriteDeadline(t time.Time) error { return nil }
 
-func setupSMTPServer(ds storage.Store) (s *Server, buf *bytes.Buffer, teardown func()) {
+func setupSMTPServer(ds storage.Store) *Server {
 	cfg := &config.Root{
 		MailboxNaming: config.FullNaming,
 		SMTP: config.SMTP{
@@ -480,26 +446,24 @@ func setupSMTPServer(ds storage.Store) (s *Server, buf *bytes.Buffer, teardown f
 			Timeout:         5,
 		},
 	}
-	// Capture log output.
-	buf = new(bytes.Buffer)
-	log.SetOutput(buf)
+
 	// Create a server, don't start it.
-	// TODO Remove teardown.
-	teardown = func() {}
 	addrPolicy := &policy.Addressing{Config: cfg}
 	manager := &message.StoreManager{Store: ds}
-	s = NewServer(cfg.SMTP, manager, addrPolicy)
-	return s, buf, teardown
+
+	return NewServer(cfg.SMTP, manager, addrPolicy)
 }
 
 var sessionNum int
 
-func setupSMTPSession(server *Server) net.Conn {
-	// Pair of pipes to communicate.
+func setupSMTPSession(t *testing.T, server *Server) net.Conn {
+	logger := zerolog.New(zerolog.NewTestWriter(t))
 	serverConn, clientConn := net.Pipe()
+
 	// Start the session.
 	server.wg.Add(1)
 	sessionNum++
-	go server.startSession(sessionNum, &mockConn{serverConn})
+	go server.startSession(sessionNum, &mockConn{serverConn}, logger)
+
 	return clientConn
 }