Paradise
Paradise

Reputation: 1468

Override a "private" method in a python module

I want to test a function in python, but it relies on a module-level "private" function, that I don't want called, but I'm having trouble overriding/mocking it. Scenario:

module.py

_cmd(command, args):
  # do something nasty

function_to_be_tested():
  # do cool things
  _cmd('rm', '-rf /')
  return 1

test_module.py

import module
test_function():
  assert module.function_to_be_tested() == 1

Ideally, in this test I dont want to call _cmd. I've looked at some other threads, and I've tried the following with no luck:

 test_function():
   def _cmd(command, args):
     # do nothing
     pass
   module._cmd = _cmd

although checking module._cmd against _cmd doesn't give the correct reference. Using mock:

from mock import patch

def _cmd_mock(command, args):
  # do nothing
  pass

@patch('module._cmd', _cmd_mock) 
test_function():
   ...

gives the correct reference when checking module._cmd, although `function_to_be_tested' still uses the original _cmd (as evidenced by it doing nasty things).

This is tricky because _cmd is a module-level function, and I dont want to move it into a module

Upvotes: 2

Views: 2350

Answers (1)

Michele d'Amico
Michele d'Amico

Reputation: 23741

[Disclaimer]

The synthetic example posted in this question works and the described issue become from specific implementation in production code. Maybe this question should be closed as off topic because the issue is not reproducible.


[Note] For impatient people Solution is at the end of the answer.


Anyway that question given to me a good point to thought: how we can patch a method reference when we cannot access to the variable where the reference is?

Lot of times I found some issue like this. There are lot of ways to meet that case and the commons are

  • Decorators: the instance we would like replace is passed as decorator argument or used in decorator static implementation
  • What we would like to patch is a default argument of a method

In both cases maybe refactor the code is the best way to play with that but what about if we are playing with some legacy code or the decorator is a third part decorator?

Ok, we have the back on the wall but we are using python and in python nothing is impossible. What we need is just the reference of the function/method to patch and instead of patching its reference we can patch the __code__: yes I'm speaking about patching the bytecode instead the function.

Get a real example. I'm using default parameter case that is simple, but it works either in decorator case.

def cmd(a):
    print("ORIG {}".format(a))
def cmd_fake(a):
    print("NEW {}".format(a))
def do_work(a, c=cmd):
    c(a)

do_work("a")
cmd=cmd_fake
do_work("b")

Output:

ORIG a
ORIG b

Ok In this case we can test do_work by passing cmd_fake but there some cases where is impossible do it: for instance what about if we need to call something like that:

def what_the_hell():
    list(map(lambda a:do_work(a), ["c","d"]))

what we can do is patch cmd.__code__ instead of _cmd by

cmd.__code__ = cmd_fake.__code__

So follow code

do_work("a")
what_the_hell()
cmd.__code__ = cmd_fake.__code__
do_work("b")
what_the_hell()

Give follow output:

ORIG a
ORIG c
ORIG d
NEW b
NEW c
NEW d

Moreover if we want to use a mock we can do it by add follow lines:

from unittest.mock import Mock, call
cmd_mock = Mock()
def cmd_mocker(a):
    cmd_mock(a)

cmd.__code__=cmd_mocker.__code__
what_the_hell()
cmd_mock.assert_has_calls([call("c"),call("d")])
print("WORKS")

That print out

WORKS

Maybe I'm done... but OP still wait for a solution of his issue

from mock import patch, Mock

cmd_mock = Mock()
#A closure for grabbing the right function code
def cmd_mocker(a):
    cmd_mock(a)

@patch.object(module._cmd,'__code__', new=cmd_mocker.__code__) 
test_function():
   ...

Now I should say never use this trick unless you are with the back on the wall. Test should be simple to understand and to debug ... try to debug something like this and you will become mad!

Upvotes: 2

Related Questions