methodtrigger.py

from threading import Lock, Condition

__all__ = ["MethodTrigger"]

class MethodTrigger:
  """
    Set a flag when a specific method is called, possibly asynchronously. Caller can wait on this flag.
    
    Example:

      class Foo(object):
        def bar(self):
          pass

      foo = Foo()
      trigger = MethodTrigger(foo, "bar")

      if trigger.wait(): # Waits for 10 seconds for foo.bar() to get called
        print "foo.bar() got called"
      else
        # This will happen, as foo.bar() wasn't called
        print "foo.bar() did not get called"

    Calls that were made before the trigger has been installed will not get recorded.
  """

  def __init__(self, obj, method):
    assert isinstance(obj, object), "Object %s does not derive from object." % (obj,)

    self.obj = obj
    self.method = method
    self.old_func = obj.__getattribute__(method)

    self.called = False
    self.args = []
    self.kwargs = {}

    self.lock = Lock()
    self.cond = Condition(self.lock)

    # Patch the target method
    obj.__setattr__(method, self.trigger)

  def trigger(self, *args, **kwargs):
    # Save the call parameters
    self.args = args
    self.kwargs = kwargs

    # Call the original method
    self.old_func(*args, **kwargs)

    # Restore the original method
    self.obj.__setattr__(self.method, self.old_func)

    # Release waiting thread
    with self.lock:
      self.called = True
      self.cond.notify()

  def wait(self, timeout=10.0):
    # Wait for method to get called
    with self.lock:
      if self.called:
        return True

      self.cond.wait(timeout)

    return self.called