Most of the code that I’ve been writing for the book has been getting its own unit tests. I’ve been working on a chapter on networking for the past week and a half and have written a little code for the chapter. One of the challenges of writing tests for something like networking code is that there are so many variables which may influence a suite of unit tests. For example, if my unit tests rely on hitting some Google webserver and I encounter problems, trouble shooting questions may include: is my router acting up, is my ISP acting up, am I failing to get DNS resolution, is that particular server down, have they changed the URL for this resource, etc.

So, for the purpose of testing, I decided to bypass the socket module in this case and handle everything locally. I created a faux socket class fleshed out with the methods that I needed. I then monkey patched my module under test with the new faux socket class. All attempts to connect to a real socket actually “connected to” a fake socket from which I could totally control the behavior.

Here’s the code to the module:

#!/usr/bin/env python

import socket
import re
import logging
logger = logging.getLogger('p4sa')

def check_webserver(address, port, resource):
    #create a TCP socket
    s = socket.socket()
    logger.info("Attempting to connect to %s on port %s" % (address, port))
    try:
        s.connect((address, port))
        logger.info("Connected to %s on port %s" % (address, port))
    except socket.error:
        logger.error("Connection to %s on port %s failed" % (address, port))
        return False
    #build up HTTP request string
    resource = re.sub('^(/)*', '/', resource)
    request_string = "GET %s HTTP/1.1nHost: %snn" % (resource, address)
    logger.info('Sending HTTP request')
    logger.debug('|||%s|||' % request_string)
    s.send(request_string)
    #we should only need the first 100 bytes or so
    rsp = s.recv(100)
    logger.info('Received 100 bytes of HTTP response')
    logger.debug('|||%s|||' % rsp)
    lines = rsp.splitlines()
    logger.info('First line of HTTP response:: %s' % lines[0])
    try:
        version, status, reason = re.split(r's+', lines[0], 2)
        logger.info('Version: %s, Status: %s, Reason: %s' % (version, status, reason))
    except ValueError:
        logger.error('Failed to split status line')
        return False
    #be a good citizen and close your connection
    s.close()
    if status in ['200']:
        logger.info('Success - status was %s' % status)
        return True
    else:
        logger.error('Status was %s' % status)
        return False

if __name__ == '__main__':
    from optparse import OptionParser
    #logger.setLevel(logging.DEBUG)
    logger.setLevel(logging.INFO)
    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)
    formatter = logging.Formatter('[%(name)s] [%(levelname)s] (%(funcName)s) %(message)s')
    ch.setFormatter(formatter)
    logger.addHandler(ch)

    parser = OptionParser()

    parser.add_option("-a", "--address", dest="address", default='localhost',
                      help="ADDRESS for webserver", metavar="ADDRESS")

    parser.add_option("-p", "--port", dest="port", type="int", default=80,
                      help="PORT for webserver", metavar="PORT")

    parser.add_option("-r", "--resource", dest="resource", default='index.html',
                      help="RESOURCE to check", metavar="RESOURCE")

    (options, args) = parser.parse_args()
    logger.debug('options: %s, args: %s' % (options, args))
    check = check_webserver(options.address, options.port, options.resource)
    logger.info('check_webserver returned %s' % check)

And here are the few unit tests I’ve written so far:

#from socket import socket as realsocket
import web_server_checker_tcp

import unittest
import logging
logger = logging.getLogger('p4sa.test')

class FauxSocketError(Exception):
    pass

class FauxSocket(object):
    def __init__(self, *args):
        #self._socket = realsocket(*args)
        logger.debug("FAUX __INIT__")
        self.resource = ""
    def connect(self, host_port):
        address, port = host_port
        if address == 'bad_connection':
            logger.error("FAUX CONNECTION FAILURE")
            raise FauxSocketError
        logger.debug("FAUX CONNECT")
        #return self._socket.connect(*args)
    def send(self, msg):
        line = msg.splitlines()[0]
        try:
            self.resource = line.split()[1]
        except IndexError:
            self.resource = ""
        logger.debug("FAUX SEND")
        #return self._socket.send(msg)
    def recv(self, bytes):
        logger.debug("FAUX RECV")
        if self.resource == '/fail_split':
            return '''ERRORnCache-Control: privatenContent-Type: text/html; charset=ISO-8859-1nSet-Cookie: PR'''
        elif self.resource == '/non_200':
            return '''HTTP/1.1 404 URL Not FoundnCache-Control: privatenContent-Type: text/html; charset=ISO-8859-1nSet-Cookie: PR'''
        else:
            return '''HTTP/1.1 200 OKnCache-Control: privatenContent-Type: text/html; charset=ISO-8859-1nSet-Cookie: PR'''
        #return self._socket.recv(bytes)
    def close(self):
        logger.debug("FAUX CLOSE")
        #return self._socket.close()

#Monkey Patching
web_server_checker_tcp.socket.socket = FauxSocket
web_server_checker_tcp.socket.error = FauxSocketError

class TestWebChecker(unittest.TestCase):
    def setUp(self):
        pass
    def testGoodResult(self):
        check = web_server_checker_tcp.check_webserver('www.google.com', 80, 'index.html')
        self.assertEqual(check, True)
    def testBadConnection(self):
        check = web_server_checker_tcp.check_webserver('bad_connection', 80, 'index.html')
        self.assertEqual(check, False)
    def testFailSplitResponse(self):
        check = web_server_checker_tcp.check_webserver('www.google.com', 80, 'fail_split')
        self.assertEqual(check, False)
    def testNon200StatusCode(self):
        check = web_server_checker_tcp.check_webserver('www.google.com', 80, 'non_200')
        self.assertEqual(check, False)

def setup_logging():
    test_logger = logging.getLogger('p4sa')
    test_logger.setLevel(logging.DEBUG)
    ch = logging.StreamHandler()
    ch.setLevel(logging.DEBUG)
    formatter = logging.Formatter('[%(name)s] [%(levelname)s] (%(funcName)s) %(message)s')
    ch.setFormatter(formatter)
    test_logger.addHandler(ch)

if __name__ == "__main__":
    setup_logging()
    unittest.main()

Basically, I created the faux socket library to know and keep track of various data that I’ve passed in and respond appropriately upon subsequent method calls. I’m sure there’s probably a better way to test this. Maybe I’ll have to replace the faux code with a mock just for the sake of doing it.

Oh, I mentioned logging. Using the logging module in Python is a great habit to get into. You can have code which spews all sorts of details about what’s going on if you have a handler to display it. When you’re running unit tests, you probably want to have the logger go silent. But if one of your unit tests starts failing all of a sudden, you can just flip on the handler and better diagnose your problem. In the “if __name__” section of the test code, I’ve been just commenting/uncommenting the “setup_logging()” line depending on whether I want to see details or not.

So, why did I write my own faux class rather than using the Python mock library? Mock objects look and feel a bit limited but convoluted to me. Especially for this case. But then again, it could be just lack of exposure to them. Maybe I need to just start using them and they won’t feel like that.

Anyway, you can expect to see code and tests something like what I’ve posted here in the book. Comments, questions, flames are all welcome.