Cleanup class decorator. Closes #363.

This commit is contained in:
Steve Pulec 2015-06-27 19:01:01 -04:00
parent 5ee1525cd8
commit 3ed9428cb0
2 changed files with 22 additions and 6 deletions

View File

@ -18,10 +18,10 @@ class MockAWS(object):
if self.__class__.nested_count == 0: if self.__class__.nested_count == 0:
HTTPretty.reset() HTTPretty.reset()
def __call__(self, func): def __call__(self, func, reset=True):
if inspect.isclass(func): if inspect.isclass(func):
return self.decorate_class(func) return self.decorate_class(func)
return self.decorate_callable(func) return self.decorate_callable(func, reset)
def __enter__(self): def __enter__(self):
self.start() self.start()
@ -64,10 +64,13 @@ class MockAWS(object):
HTTPretty.disable() HTTPretty.disable()
HTTPretty.reset() HTTPretty.reset()
def decorate_callable(self, func): def decorate_callable(self, func, reset):
def wrapper(*args, **kwargs): def wrapper(*args, **kwargs):
with self: self.start(reset=reset)
try:
result = func(*args, **kwargs) result = func(*args, **kwargs)
finally:
self.stop()
return result return result
functools.update_wrapper(wrapper, func) functools.update_wrapper(wrapper, func)
wrapper.__wrapped__ = func wrapper.__wrapped__ = func
@ -87,7 +90,7 @@ class MockAWS(object):
continue continue
try: try:
setattr(klass, attr, self(attr_value)) setattr(klass, attr, self(attr_value, reset=False))
except TypeError: except TypeError:
# Sometimes we can't set this for built-in types # Sometimes we can't set this for built-in types
continue continue

View File

@ -2,10 +2,12 @@ from __future__ import unicode_literals
import boto import boto
from boto.exception import EC2ResponseError from boto.exception import EC2ResponseError
import sure # noqa import sure # noqa
import unittest
import tests.backport_assert_raises # noqa import tests.backport_assert_raises # noqa
from nose.tools import assert_raises from nose.tools import assert_raises
from moto import mock_ec2 from moto import mock_ec2, mock_s3
''' '''
Test the different ways that the decorator can be used Test the different ways that the decorator can be used
@ -69,3 +71,14 @@ class Tester(object):
def test_still_the_same(self): def test_still_the_same(self):
conn = boto.connect_ec2() conn = boto.connect_ec2()
list(conn.get_all_instances()).should.have.length_of(0) list(conn.get_all_instances()).should.have.length_of(0)
@mock_s3
class TesterWithSetup(unittest.TestCase):
def setUp(self):
self.conn = boto.connect_s3()
self.conn.create_bucket('mybucket')
def test_still_the_same(self):
bucket = self.conn.get_bucket('mybucket')
bucket.name.should.equal("mybucket")