djangoproject.com | nginx.org | python.org | linux.com
version seven.
  http://demongin.org
demongin.org - Python Decorators: Wrapping Test Functions Without Losing Methods

Python Decorators: Wrapping Test Functions Without Losing Methods

A kind of insane-sounding how-to describing a Selenium object created by a unittest function that gets wrapped by a decorator that calls a method from the Selenium class: the result is a Selenium test that can kill its browser if the test raises an Exception. Not for the faint of heart.


Monday, 2010-01-18 | Careerism, Programming, Testing

It is wise to keep in mind that neither success nor failure is ever final.
Roger Babson

I have blogged about preserving self while wrapping python functions before. In this old post, I describe using the wraps built-in to pass self to the decorator (and do whatever you want with it.

Recently I've been getting up to some new chicanery with decorators. The latest death- and common-sense-defying feat of decoration I've dreamed up involves calling methods of class objects created by functions.

If it sounds goofy, that's because it kind of is. But bear with me, for a second, and you'll see how this sort of thing can be useful, especially when writing functions that are going to be part of an iterative process, e.g. in the case of a test suite (or any case where you've got functions that you want to behave a certain way if they're decorated with one thing or a different way if they're decorated with another thing).

So, first things first, here's a short description of what I was doing and what I eventually realized that I wanted to do.

When I started, I had a unittest class where each function of the class was a test. This is SOP when writing python unit tests: you make a class, make each function a test (that ends with some kind of assert statement and then execute the main() function of the class when you want to run them all in order. Here's what my unittest class looked like: 

class SortClients(unittest.TestCase):
    sel = Browser()
    sel.setUp()
    active_browser = sel.log_in()

    def test_reverse_sort(self): 
        s = self.sel.sort_test(self.active_browser)
        self.assert_(s == True, "Reverse sort failed!") 

if __name__ == "__main__":
    unittest.main()
Now, if you looked carefully at that, you'll have noticed a few things:
  1. In the first line of my class, I instantiate another class called Browser() and call the instance sel. This, you may have already guessed, is a Selenium object.
  2. Once I've got sel, I call its setUp() method and then create a new Selenium object called active_browser. This is exactly what it sounds like: it's an open browser window on my Selenium server that's waiting for me to throw instructions at it.
  3. The first (and only) test in my class is kind of vanilla: it calls the sort_test method of sel and returns the results, which are then asserted and used to determine the outcome of the test.
Now, the problem I was having--the problem that I am about to solve with decorators--was that if something unexpected went wrong when I called sort_test, then the test would determine that its result was "Exception" (as opposed to pass or fail) and leave the Selenium browser open and stalled out on the test server.

And honestly, this wasn't a huge problem... at first. For about two months, things went on like this: if a test had an exception, it would just leave the window open and the Firefox process running. But, as idle processes tend to do, these FF processes stacked up and eventually caused the server to become sluggish and, eventually, unresponsive: when I finally popped open a VNC window on the virtual desktop where the tests were being run, there were so many of them that they had broken Gnome's task bar--I eventually wound up tossing off a pkill firefox and took down something like 10 or 11 thousand PID's.

And so it became clear to me that what I wanted to do was decorate certain tests (i.e. functions) in my test cases (i.e. classes) so that if those certain tests caused an Exception to be raised, the Selenium-controlled browser would be killed. I didn't want this on all tests, of course, because some of my tests depended on previous tests having already opened the browser. Where I wanted this functionality was usually on the last test (i.e. function) in my test case.

And, as I rapidly learned, I couldn't just write some code into the unittest functions. As soon as they raise an Exception, that's it for them: they do don't do try/except, they don't wait around for any subsequent code, they don't pass "GO" and they sure as Hell don't collect $200--they simply return the Exception to the main() function and die quietly. Use of a decorator that could detect an Exception and call a different function before the test function could raise an exception was clearly indicated.

More specifically, I realized that I needed a decorator that took the function it decorated, got self, got an object "within" self and then, finally, called the tearDown() method of that object. Here's what I came up with:
def tear_down(func):
    @wraps(func)
    def new(*args, **kwargs):
        try:
            func(*args, **kwargs)
            args[0].tearDown.__self__.sel.tearDown()
        except Exception, e:
            args[0].tearDown.__self__.sel.tearDown()
            raise
    return new
As you can see, the decorator uses the wraps built-in discussed above, to preserve self. Then it uses a basic try/except routine to either a.) run the tearDown() function when the decorated function does not raise an Exception or b.) run the tearDown() function when the decorated function does raise an exception and then raise that exception, effectively ending the test and reporting the result as an Exception (which the unittest class then reports in its results.

So yeah, like I said, this is kind of esoteric/obscure python, but definitely also the sort of thing that can help you slim down your codebase and add some choice functionality to your iterative scripting.