blob: 18ee0cb70768ed3a83dbb440fbf4209e9eda0250 [file] [log] [blame]
# Copyright 2015 Google Inc. All Rights Reserved.
"""Utilities for unit tests."""
class ReplaceFunction(object):
"""Makes it easier to replace a function in a class or module."""
def __init__(self, obj, function_name, new_function):
self.obj = obj
self.function_name = function_name
self.old_function = getattr(self.obj, self.function_name)
setattr(self.obj, self.function_name, new_function)
def __del__(self):
setattr(self.obj, self.function_name, self.old_function)
class CapturingFunction(object):
"""Captures the last arguments called on a function."""
def __init__(self, return_value=None):
self.num_calls = 0
self.return_value = return_value
self.last_args = None
self.last_kwargs = None
def __call__(self, *args, **kwargs):
self.last_args = args
self.last_kwargs = kwargs
self.num_calls += 1
if callable(self.return_value):
return self.return_value()
return self.return_value