1 files changed, 18 insertions, 27 deletions

diff --git a/src/mailman/testing/mta.py b/src/mailman/testing/mta.py
index e89cc7cfb..28e5ec77c 100644
--- a/src/mailman/testing/mta.py
+++ b/src/mailman/testing/mta.py
@@ -59,20 +59,6 @@ class FakeMTA:
 class SessionCountingChannel(Channel):
     """Count the number of SMTP sessions opened and closed."""
 
-    def smtp_HELO(self, arg):
-        """See `smtpd.SMTPChannel.smtp_HELO`."""
-        # Store this on the server because while the channel has access to the
-        # server, the server does not have access to the individual channels.
-        self._server.helo_count += 1
-        Channel.smtp_HELO(self, arg)
-
-    def smtp_QUIT(self, arg):
-        """See `smtpd.SMTPChannel.smtp_QUIT`."""
-        # Store this on the server because while the channel has access to the
-        # server, the server does not have access to the individual channels.
-        self._server.quit_count += 1
-        Channel.smtp_QUIT(self, arg)
-
     def smtp_STAT(self, arg):
         """Cause the server to send statistics to its controller."""
         self._server.send_statistics()
@@ -86,11 +72,7 @@ class SessionCountingServer(QueueServer):
     def __init__(self, host, port, queue, oob_queue):
         """See `lazr.smtptest.server.QueueServer`."""
         QueueServer.__init__(self, host, port, queue)
-        # Store these on the server because while the channel has access to
-        # the server, the server does not have access to the individual
-        # channels.
-        self.helo_count = 0
-        self.quit_count = 0
+        self.session_count = 0
         # The out-of-band queue is where the server sends statistics to the
         # controller upon request.
         self._oob_queue = oob_queue
@@ -98,18 +80,20 @@ class SessionCountingServer(QueueServer):
     def handle_accept(self):
         """See `lazr.smtp.server.Server`."""
         connection, address = self.accept()
+        self.session_count += 1
         log.info('[SessionCountingServer] accepted: %s', address)
         SessionCountingChannel(self, connection, address)
 
     def reset(self):
         """See `lazr.smtp.server.Server`."""
         QueueServer.reset(self)
-        self.helo_count = 0
-        self.quit_count = 0
+        self.session_count = 0
 
     def send_statistics(self):
         """Send the current connection statistics to the controller."""
-        self._oob_queue.put((self.helo_count, self.quit_count))
+        # Do not count the connection caused by the STAT connect.
+        self.session_count -= 1
+        self._oob_queue.put(self.session_count)
 
 
 
@@ -126,17 +110,24 @@ class SessionCountingController(QueueController):
         self.server = SessionCountingServer(
             host, port, self.queue, self.oob_queue)
 
-    def get_statistics(self):
-        """Retrieve connection statistics from the server.
+    def start(self):
+        """See `lazr.smtptest.controller.QueueController`."""
+        QueueController.start(self)
+        # Reset the connection statistics, since the base class's start()
+        # method causes a connection to occur.
+        self.reset()
+
+    def get_session_count(self):
+        """Retrieve the number of sessions.
 
-        :return: a 2-tuple of the format (HELO count, QUIT count)
-        :rtype 2-tuple of integers
+        :return: The number of sessions that have been opened.
+        :rtype: integer
         """
         smtpd = self._connect()
         smtpd.docmd('STAT')
         # An Empty exception will occur if the data isn't available in 10
         # seconds.  Let that propagate.
-        return self.queue.get(block=True, timeout=10)
+        return self.oob_queue.get(block=True, timeout=10)
 
     @property
     def messages(self):