Refactor SMTP handler_test log collection (#289)

* smtp: allow logger to be passed into startSession
* smtp: remove logbuf and teardown from handler_test
* smtp: handler_test log output to testing.T

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) {
+func (s *Server) startSession(id int, conn net.Conn, logger zerolog.Logger) {
-	logger := log.Hook(logHook{}).With().
+	logger = logger.Hook(logHook{}).With().
 		Str("module", "smtp").
 		Str("remote", conn.RemoteAddr().String()).
 		Int("session", id).Logger()

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)
+	server := setupSMTPServer(ds)
-	defer teardown()
+	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)
+	server := setupSMTPServer(ds)
-	defer teardown()
+	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)
+	server := setupSMTPServer(ds)
-	defer teardown()
+	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)
+	server := setupSMTPServer(ds)
-	defer teardown()
+	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)
+	server := setupSMTPServer(mds)
-	defer teardown()
+	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)
+	server := setupSMTPServer(mds)
-	defer teardown()
+	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,6 +365,7 @@ Hi!
 	if err := playScriptAgainst(t, c, script); err != nil {
 		t.Error(err)
 	}
+
 	// Send a message
 	body = `X-Useless-Header: true
 
@@ -406,18 +377,13 @@ Hi! Can you still deliver this?
 	if code, _, err := c.ReadCodeLine(250); err != nil {
 		t.Errorf("Expected a 250 greeting, got %v", code)
 	}
-
+	_, _ = c.Cmd("QUIT")
-	if t.Failed() {
+	_, _, _ = c.ReadCodeLine(221)
-		// 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)
-	}
 }
 
 // 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 {
+func setupSMTPSession(t *testing.T, server *Server) net.Conn {
-	// Pair of pipes to communicate.
+	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
 }

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)
 		}
 	}
 }