adrift on a cosmic ocean

Writings on various topics (mostly technical) from Oliver Hookins and Angela Collins. We have lived in Berlin since 2009, have two kids, and have far too little time to really justify having a blog.

Mocking fun with Python, mox and socket

Posted by Oliver on the 9th of September, 2011 in category Tech
Tagged with: moxpymoxpythonsocketunittest

I've been doing most of my work with Ruby this last year, as a lot of it was Puppet-related. I went through the pain and pleasure of learning Ruby and the intricacies of unit testing with test::unit and RSpec, and mocking objects with the excellent Mocha library. Now I have slightly returned to my roots by doing a bit of Python programming on another part of our codebase.

We use Python for some of the code we use to test our infrastructure - initially because it most suited the engineers who were implementing it and Python was already on our list of supported languages. We don't want our language set to grow unwieldy but at this point in time we have a fairly well-defined set of domains in which we restrict our language choice. Finally in this sprint my chance came up to do some work on this code base so it was time to freshen my memory on Python unit testing.

The particular task at hand was verifying connectivity to the management interface of a piece of hardware. Without giving too much away, you might for example do an extremely simple check of a device that has an SSH interface as follows:

    # Create a TCP connection to the SSH port and grab the output
    try:
      try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(5)
        sock.connect((ip, socket.getservbyname('ssh')))
        sock_output = sock.recv(100)
      except socket.timeout:
        return False, "SSH Connection Timeout"
      except socket.error:
        return False, "SSH Connection Error"
    finally:
      sock.close()

    result = re.search('SSH-2.0-OpenSSH_5.1', sock_output)
    if not result:
      return False, "Valid SSH interface not found"
    else:
      return True, "Valid SSH interface was found"

No, this is not great code (and you can blame Python 2.4 for that crazy nested exception handling block), but it works and reasonably catches the obvious errors. How do we go about testing this? I have another test class which verifies connectivity to an HTTP interface of a piece of hardware, but the decision was just about made up for me with that. Since the interface actually requires SSL, it would be quite difficult, tiring and time consuming to make a real HTTPS connection to a fake server and pass some dummy text across the connection. Therefore in that case I simply stubbed the calls to urllib2.urlopen and passed the dummy text without any network being involved.

In the case of a much simpler interface like a socket, we have a quandary on our hands:

  • Create a listening socket, have it call a real object which passes data over the connection which goes back to the connecting object.
  • Create a listening socket which mocks out the send method in order to send some pre-cooked output back to the connecting object.
  • Mock the sending socket and have pre-cooked output going back to the calling object.

Ideally we'd keep things as real as possible, spinning up a connected pair of sockets on the loopback adaptor and performing reasonably real communication, but after being away from Python for a while it was far too taxing on the brain. To cut a long story short I decided to take the last option and mock the socket object itself, although as it turned out, it was a bit tricky:

  def setUp(self):
    """ Create a mock socket for later use """
    super(TestSshTest, self).setUp()
    self.mock_socket = self.mox.CreateMockAnything(socket.socket)
    self.mox.StubOutWithMock(socket, 'getservbyname')
    self.mox.StubOutWithMock(socket, 'socket')
    socket.socket(socket.AF_INET, socket.SOCK_STREAM).AndReturn(self.mock_socket)

  def test_ssh_success(self):
    """ Standard successful test """
    socket.getservbyname('ssh').AndReturn(33333)
    self.mock_socket.connect(('192.0.2.1', 33333))
    self.mock_socket.recv(100).AndReturn('SSH-2.0-OpenSSH_5.1')
    self.mock_socket.close()
    self.mox.ReplayAll()

    # real method runs here

Don't ask me exactly why all the above was necessary, as I'm still coming to grips with it. Part of it I believe is due to the intricacies of exactly how different entities work in Python, and part of it is how sockets are set up. Pymox appears to be more complicated to use than Mocha but I'm reasonably certain that is just added flexibility and power that Mocha hides from you - I'm still waiting for the lightbulb moment.

If anybody out there has successfully mocked out sockets in Python using Pymox and in a way that is simpler than this, I'd love to hear how it was done!

© 2010-2018 Oliver Hookins and Angela Collins