Flexmock with decorators test

gistfile1.py

#!/usr/bin/python

import functools
import unittest

from flexmock.flexmock import flexmock


def FooDecorator(fn):

  @functools.wraps(fn)
  def Wrapper(*args, **kwargs):
    return fn(*args, **kwargs)

  return Wrapper


def ComplicatedDecorator(arbitrary_name):

  def Decorator(fn):

    @functools.wraps(fn)
    def Wrapper(*args, **kwargs):
      return arbitrary_name

    return Wrapper

  return Decorator


@FooDecorator
def SomeFunction(*args, **kwargs):
  return True


@ComplicatedDecorator('arbitrary name')
def OtherFunction(*args, **kwargs):
  return True


class FooClass(object):

  @FooDecorator
  def SomeMethod(self, *args, **kwargs):
    return True

  @ComplicatedDecorator('arbitrary name')
  def OtherMethod(self, *args, **kwargs):
    return True


class DecoratorTest(unittest.TestCase):

  def testFlexmockDecoratedFunction(self):
    flexmock().should_receive('SomeFunction').and_return('foo').once
    self.assertEqual('foo', SomeFunction())

  def testFlexmockComplicatedDecoratedFunction(self):
    flexmock().should_receive('OtherFunction').and_return('foo').once
    self.assertEqual('foo', OtherFunction())

  def testFlexmockDecoratorMethod(self):
    flexmock(FooClass).should_receive('SomeMethod').and_return('foo').once
    self.assertEqual('foo', FooClass.SomeMethod())

  def testFlexmockComplicatedDecoratedMethod(self):
    flexmock(FooClass).should_receive('OtherMethod').and_return('foo').once
    self.assertEqual('foo', FooClass.OtherMethod())

  def testMySanity(self):

    def Replacement():
      return 'foo'

    try:
      old_other_function = OtherFunction
      globals()['OtherFunction'] = Replacement
      self.assertEqual(Replacement, OtherFunction)
      self.assertEqual('foo', OtherFunction())
    finally:
      globals()['OtherFunction'] = old_other_function



if __name__ == '__main__':
  unittest.main()