18

I have decorator @login_testuser applied to method test_1():

class TestCase(object):
    @login_testuser
    def test_1(self):
        print "test_1()"

Is there a way I can apply @login_testuser on every method of the class prefixed with "test_"?

In other words, the decorator would apply to test_1(), test_2() methods below, but not on setUp().

class TestCase(object):
    def setUp(self):
        pass

    def test_1(self):
        print "test_1()"

    def test_2(self):
        print "test_2()"
martineau
  • 119,623
  • 25
  • 170
  • 301
Jeff Bauer
  • 13,890
  • 9
  • 51
  • 73

4 Answers4

26

In Python 2.6, a class decorator is definitely the way to go. e.g., here's a pretty general one for these kind of tasks:

import inspect

def decallmethods(decorator, prefix='test_'):
    def dectheclass(cls):
        for name, m in inspect.getmembers(cls, inspect.isfunction):
            if name.startswith(prefix):
                setattr(cls, name, decorator(m))
        return cls
    return dectheclass


@decallmethods(login_testuser)
class TestCase(object):
    def setUp(self):
        pass

    def test_1(self):
        print("test_1()")

    def test_2(self):
        print("test_2()")

will get you what you desire. In Python 2.5 or worse, the @decallmethods syntax doesn't work for class decoration, but with otherwise exactly the same code you can replace it with the following statement right after the end of the class TestCase statement:

TestCase = decallmethods(login_testuser)(TestCase)
martineau
  • 119,623
  • 25
  • 170
  • 301
Alex Martelli
  • 854,459
  • 170
  • 1,222
  • 1,395
5

Sure. Iterate all attributes of the class. Check each one for being a method and if the name starts with "test_". Then replace it with the function returned from your decorator

Something like:

from inspect import ismethod, getmembers
for name, obj in getmembers(TestCase, ismethod):
   if name.startswith("test_"):
       setattr(TestCase, name, login_testuser(obj))
mthurlin
  • 26,247
  • 4
  • 39
  • 46
2

Are you sure you wouldn't be better off by putting login_testuser's code into setUp instead? That's what setUp is for: it's run before every test method.

Ned Batchelder
  • 364,293
  • 75
  • 561
  • 662
  • Thanks for the suggestion, but I wish to apply the decorator only to methods with a particular prefix, and setUp will apply to all methods. – Jeff Bauer Feb 10 '10 at 17:02
  • 2
    Perhaps segregate those methods into their own class then? Magic inspection and augmentation of methods sounds like a mysterious maintenance headache. – Ned Batchelder Feb 10 '10 at 18:48
0

Yes, you can loop over the class's dir/__dict__ or have a metaclass that does so, identifying if the attributes start with "test". However, this will create less straightforward, explicit code than writing the decorator explicitly.

Mike Graham
  • 73,987
  • 14
  • 101
  • 130