pygrader/test/server.py

   1 # Copyright (C) 2012 W. Trevor King <wking@tremily.us>
   2 #
   3 # This file is part of pygrader.
   4 #
   5 # pygrader is free software: you can redistribute it and/or modify it under the
   6 # terms of the GNU General Public License as published by the Free Software
   7 # Foundation, either version 3 of the License, or (at your option) any later
   8 # version.
   9 #
  10 # pygrader is distributed in the hope that it will be useful, but WITHOUT ANY
  11 # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
  12 # A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
  13 #
  14 # You should have received a copy of the GNU General Public License along with
  15 # pygrader.  If not, see <http://www.gnu.org/licenses/>.
  16
  17 import asyncore as _asyncore
  18 import email as _email
  19 import smtpd as _smptd
  20 import socket as _socket
  21
  22 from .. import LOG as _LOG
  23
  24
  25 class SMTPChannel (_smptd.SMTPChannel):
  26     def close(self):
  27         super(SMTPChannel, self).close()
  28         _LOG.debug('close {}'.format(self))
  29         self.smtp_server.channel_closed()
  30
  31
  32 class SMTPServer (_smptd.SMTPServer):
  33     """An SMTP server for testing pygrader.
  34
  35     >>> from asyncore import loop
  36     >>> from smtplib import SMTP
  37     >>> from pgp_mime.email import encodedMIMEText
  38     >>> from pygrader.test.client import MessageSender
  39
  40     >>> def process(peer, mailfrom, rcpttos, data):
  41     ...     print('peer:     {}'.format(peer))
  42     ...     print('mailfrom: {}'.format(mailfrom))
  43     ...     print('rcpttos:  {}'.format(rcpttos))
  44     ...     print('message:')
  45     ...     print(data)
  46     >>> server = SMTPServer(
  47     ...     ('localhost', 1025), None, process=process, count=3)
  48
  49     >>> message = encodedMIMEText('Ping')
  50     >>> message['From'] = 'a@example.com'
  51     >>> message['To'] = 'b@example.com, c@example.com'
  52     >>> message['Cc'] = 'd@example.com'
  53     >>> messages = [message, message, message]
  54     >>> ms = MessageSender(address=('localhost', 1025), messages=messages)
  55     >>> loop()  # doctest: +REPORT_UDIFF, +ELLIPSIS
  56     peer:     ('127.0.0.1', ...)
  57     mailfrom: a@example.com
  58     rcpttos:  ['b@example.com', 'c@example.com', 'd@example.com']
  59     message:
  60     Content-Type: text/plain; charset="us-ascii"
  61     MIME-Version: 1.0
  62     Content-Transfer-Encoding: 7bit
  63     Content-Disposition: inline
  64     From: a@example.com
  65     To: b@example.com, c@example.com
  66     Cc: d@example.com
  67     <BLANKLINE>
  68     Ping
  69     peer:     ('127.0.0.1', ...)
  70     mailfrom: a@example.com
  71     rcpttos:  ['b@example.com', 'c@example.com', 'd@example.com']
  72     message:
  73     Content-Type: text/plain; charset="us-ascii"
  74     MIME-Version: 1.0
  75     Content-Transfer-Encoding: 7bit
  76     Content-Disposition: inline
  77     From: a@example.com
  78     To: b@example.com, c@example.com
  79     Cc: d@example.com
  80     <BLANKLINE>
  81     Ping
  82     peer:     ('127.0.0.1', ...)
  83     mailfrom: a@example.com
  84     rcpttos:  ['b@example.com', 'c@example.com', 'd@example.com']
  85     message:
  86     Content-Type: text/plain; charset="us-ascii"
  87     MIME-Version: 1.0
  88     Content-Transfer-Encoding: 7bit
  89     Content-Disposition: inline
  90     From: a@example.com
  91     To: b@example.com, c@example.com
  92     Cc: d@example.com
  93     <BLANKLINE>
  94     Ping
  95     """
  96     channel_class = SMTPChannel
  97
  98     def __init__(self, *args, **kwargs):
  99         self.count = kwargs.pop('count', None)
 100         self.process = kwargs.pop('process', None)
 101         self.channels_open = 0
 102         super(SMTPServer, self).__init__(*args, **kwargs)
 103
 104     def log_info(self, message, type='info'):
 105         # TODO: type -> severity
 106         _LOG.info(message)
 107
 108     def handle_accepted(self, conn, addr):
 109         if self.count <= 0:
 110             conn.close()
 111             return
 112         super(SMTPServer, self).handle_accepted(conn, addr)
 113         self.channels_open += 1
 114
 115     def channel_closed(self):
 116         self.channels_open -= 1
 117         if self.channels_open == 0 and self.count <= 0:
 118             _LOG.debug('close {}'.format(self))
 119             self.close()
 120
 121     def process_message(self, peer, mailfrom, rcpttos, data):
 122         if self.count is not None:
 123             self.count -= 1
 124             _LOG.debug('Count: {}'.format(self.count))
 125         _LOG.debug('receiving message from: {}'.format(peer))
 126         _LOG.debug('message addressed from: {}'.format(mailfrom))
 127         _LOG.debug('message addressed to  : {}'.format(rcpttos))
 128         _LOG.debug('message length        : {}'.format(len(data)))
 129         if self.process:
 130             self.process(peer, mailfrom, rcpttos, data)
 131         return