add in smtp connections

README.rst

@@ -9,12 +9,21 @@ Then add the following to your tests::
 
   --with-blockage
 
+Blocking HTTP
+-------------
+
 By default it whitelists `localhost` and `127.0.0.1`. To change the whitelist::
 
   --http-whitelist=some.site,some.other.site
 
 If the code hits a http connection then instead of completing it will raise a
 MockHTTPCall exception. Please go and mock your tests appropriately.
 
-To come, maybe more libs.
+Blocking SMTP
+-------------
+
+By default it whitelists no domains. To change the whitelist::
+
+  --smtp-whitelist=some.site
 
+It will raise a MockSMTPCall exception.

blockage/plugins.py

@@ -1,6 +1,7 @@
 import httplib
 import logging
 import os
+import smtplib
 
 from nose.plugins.base import Plugin
 
@@ -12,26 +13,43 @@ class MockHTTPCall(Exception):
     # for you so you'll remember to do it next time.
     pass
 
+class MockSMTPCall(Exception):
+    # If you don't mock out smtp calls in your tests, we'll raise an error
+    # for you so you'll remember to do it next time.
+    pass
+
 
 class NoseBlockage(Plugin):
     name = 'blockage'
     default_http_whitelist = '127.0.0.1,localhost'
+    default_smtp_whitelist = ''
 
     def options(self, parser, env=os.environ):
         super(NoseBlockage, self).options(parser, env=env)
         parser.add_option('--http-whitelist', action='store',
                           default=env.get('HTTP_WHITELIST',
                                           self.default_http_whitelist),
                           dest='http_whitelist')
+        parser.add_option('--smtp-whitelist', action='store',
+                          default=env.get('SMTP_WHITELIST',
+                                          self.default_smtp_whitelist),
+                          dest='smtp_whitelist')
 
     def configure(self, options, conf):
         self.options = options
         super(NoseBlockage, self).configure(options, conf)
-        self.http_whitelist = [s.strip() for s in
-                               self.options.http_whitelist.split(',')]
+        self.whitelists = {}
+        for name, option in (['http', 'http_whitelist'],
+                             ['smtp', 'smtp_whitelist']):
+            self.whitelists[name] = [s.strip() for s in
+                                     getattr(self.options, option).split(',')]
 
     def begin(self):
-        http_whitelist = self.http_whitelist
+        self.do_http()
+        self.do_smtp()
+
+    def do_http(self):
+        http_whitelist = self.whitelists['http']
 
         def whitelisted(self, host, *args, **kwargs):
             if isinstance(host, basestring) and host not in http_whitelist:
@@ -46,3 +64,20 @@ def whitelisted(self, host, *args, **kwargs):
             log.debug('Monkey patching httplib')
             httplib.HTTPConnection.old = httplib.HTTPConnection.__init__
             httplib.HTTPConnection.__init__ = whitelisted
+
+    def do_smtp(self):
+        smtp_whitelist = self.whitelists['smtp']
+
+        def whitelisted(self, host, *args, **kwargs):
+            if isinstance(host, basestring) and host not in smtp_whitelist:
+                log.warning('Denied SMTP connection to: %s' % host)
+                raise MockSMTPCall(host)
+            log.debug('Allowed SMTP connection to: %s' % host)
+            return self.old(host, *args, **kwargs)
+
+        whitelisted.blockage = True
+
+        if not getattr(httplib.HTTPConnection, 'blockage', False):
+            log.debug('Monkey patching smtplib')
+            smtplib.SMTP.old = smtplib.SMTP.__init__
+            smtplib.SMTP.__init__ = whitelisted