Python patch existing class

Question

I'm an experienced PHP/Ruby developer but right now I'm fighting Python and I really need your help.

I need to patch existing class by adding static attribute and overwriting static function to use it.

Let me show you example:

class Test():
    @staticmethod
    def generate():
        return 10

But in my test suite I need to get the following class:

class Test():
    count = 1
    @staticmethod
    def generate():
        if Test.count < 3:
            Test.count += 1
            return 1
        else:
            return 10

So the basic idea is to get 10 only on the 3rd call of "generate" function.

My first approach was to use "patch" technique, so I did:

def my_generate_hash():
    return 99

with patch.object(Test, 'generate', staticmethod(my_generate_hash)):
    print "Got %d" % check_hash()

Buuut I was unable to implement attribute "count" and use it in overriding method (

Second thought was to "Mock" something! So..

mock = MagicMock(Test)
mock.count = 1
def my_generate_hash():
    if Test2.count < 3:
        Test2.count += 1
        return 1
    else:
        return 10
mock.generate = my_generate_hash
with patch('__main__.Test', mock):
    print Test.generate()

But in real case I have other methods in "Test" class, so it won't work.

I'm stuck. Any help will be appreciated!

Answer 1

Looks like in can be in a different way.

self.count = 0

def generate():
if self.count < 3 
     self.count += 1
     return 10
else:         
    return 99

with patch.object(Test, 'generate', generate):
    self.assertEqual(Test.generate(), 10)
    self.assertEqual(Test.generate(), 10)
    self.assertEqual(Test.generate(), 10)
    self.assertEqual(Test.generate(), 99)
    self.assertEqual(Test.generate(), 99)

Answer 2

It might be simpler to subclass the original Test class for use in your tests:

class Test(object):
    @staticmethod
    def generate():
        return 10

class PatchedTest(Test):
    count = 1
    @staticmethod
    def generate():
        if Test.count < 3:
            Test.count += 1
            return 1
        else:
            return 10

The replacement function could also be done in two somewhat better ways, both of which should make it a lot easier to patch the Test class in the same way you were trying in your question:

Use a @classmethod, allowing the function to access the class it's assigned to:

class PatchedTest(Test):
    count = 1

    @classmethod
    def generate(cls):
        if cls.count < 3:
            cls.count += 1
            return 1
        else:
            return 10

Use a generator instead - each time the function is called it will continue execution where it last left off. However, this will only work if you are iterating over the functions result:

 def alternative_generate():
     yield 1
     yield 1
     yield 10