From 7816b191b87d7d69e100d68f45d1ee05551861bc Mon Sep 17 00:00:00 2001 From: Bert Blommers Date: Wed, 22 Feb 2023 10:21:45 -0100 Subject: [PATCH] Core: Reset data when mock stops (#5958) --- moto/core/models.py | 24 +++++++--- tests/test_core/test_moto_api.py | 63 ++++++++++++++++++++++++++- tests/test_s3/test_s3_file_handles.py | 16 +++++++ 3 files changed, 94 insertions(+), 9 deletions(-) diff --git a/moto/core/models.py b/moto/core/models.py index 68497694c..58c3b24f3 100644 --- a/moto/core/models.py +++ b/moto/core/models.py @@ -64,10 +64,12 @@ class BaseMockAWS: if self.__class__.nested_count == 0: self.reset() # type: ignore[attr-defined] - def __call__(self, func: Callable[..., Any], reset: bool = True) -> Any: + def __call__( + self, func: Callable[..., Any], reset: bool = True, remove_data: bool = True + ) -> Any: if inspect.isclass(func): return self.decorate_class(func) - return self.decorate_callable(func, reset) + return self.decorate_callable(func, reset, remove_data) def __enter__(self) -> "BaseMockAWS": self.start() @@ -89,7 +91,7 @@ class BaseMockAWS: self.enable_patching(reset) # type: ignore[attr-defined] - def stop(self) -> None: + def stop(self, remove_data: bool = True) -> None: self.__class__.nested_count -= 1 if self.__class__.nested_count < 0: @@ -105,18 +107,23 @@ class BaseMockAWS: pass self.unmock_env_variables() self.__class__.mocks_active = False - reset_model_data() + if remove_data: + # Reset the data across all backends + for backend in self.backends.values(): + backend.reset() + # Remove references to all model instances that were created + reset_model_data() self.disable_patching() # type: ignore[attr-defined] def decorate_callable( - self, func: Callable[..., CALLABLE_RETURN], reset: bool + self, func: Callable[..., CALLABLE_RETURN], reset: bool, remove_data: bool ) -> Callable[..., CALLABLE_RETURN]: def wrapper(*args: Any, **kwargs: Any) -> CALLABLE_RETURN: self.start(reset=reset) try: result = func(*args, **kwargs) finally: - self.stop() + self.stop(remove_data=remove_data) return result functools.update_wrapper(wrapper, func) @@ -173,15 +180,18 @@ class BaseMockAWS: # Special case for UnitTests-class is_test_method = attr.startswith(unittest.TestLoader.testMethodPrefix) should_reset = False + should_remove_data = False if attr in ["setUp", "setup_method"]: should_reset = True elif not has_setup_method and is_test_method: should_reset = True + should_remove_data = True else: # Method is unrelated to the test setup # Method is a test, but was already reset while executing the setUp-method pass - setattr(klass, attr, self(attr_value, reset=should_reset)) + kwargs = {"reset": should_reset, "remove_data": should_remove_data} + setattr(klass, attr, self(attr_value, **kwargs)) except TypeError: # Sometimes we can't set this for built-in types continue diff --git a/tests/test_core/test_moto_api.py b/tests/test_core/test_moto_api.py index 3d69ba23b..cac92f034 100644 --- a/tests/test_core/test_moto_api.py +++ b/tests/test_core/test_moto_api.py @@ -5,8 +5,9 @@ import boto3 import json import pytest from botocore.exceptions import ClientError -from moto import mock_autoscaling, mock_sqs, settings -from unittest import SkipTest +from moto import mock_autoscaling, mock_s3, mock_sqs, settings +from moto.core.model_instances import model_data, reset_model_data +from unittest import SkipTest, TestCase base_url = ( "http://localhost:5000" @@ -73,3 +74,61 @@ def test_creation_error__data_api_still_returns_thing(): names = [obj["name"] for obj in as_objects["FakeAutoScalingGroup"]] names.should.contain("test_asg") + + +def test_model_data_is_emptied_as_necessary(): + if settings.TEST_SERVER_MODE: + raise SkipTest("We're only interested in the decorator performance here") + + # Reset any residual data + reset_model_data() + + # No instances exist, because we have just reset it + for classes_per_service in model_data.values(): + for _class in classes_per_service.values(): + _class.instances.should.equal([]) + + with mock_sqs(): + # When just starting a mock, it is empty + for classes_per_service in model_data.values(): + for _class in classes_per_service.values(): + _class.instances.should.equal([]) + + # After creating a queue, some data will be present + conn = boto3.client("sqs", region_name="us-west-1") + conn.create_queue(QueueName="queue1") + + model_data["sqs"]["Queue"].instances.should.have.length_of(1) + + # But after the mock ends, it is empty again + for classes_per_service in model_data.values(): + for _class in classes_per_service.values(): + _class.instances.should.equal([]) + + # When we have multiple/nested mocks, the data should still be present after the first mock ends + with mock_sqs(): + conn = boto3.client("sqs", region_name="us-west-1") + conn.create_queue(QueueName="queue1") + with mock_s3(): + # The data should still be here - instances should not reset if another mock is still active + model_data["sqs"]["Queue"].instances.should.have.length_of(1) + # The data should still be here - the inner mock has exited, but the outer mock is still active + model_data["sqs"]["Queue"].instances.should.have.length_of(1) + + +@mock_sqs +class TestModelDataResetForClassDecorator(TestCase): + def setUp(self): + if settings.TEST_SERVER_MODE: + raise SkipTest("We're only interested in the decorator performance here") + + # No data is present at the beginning + for classes_per_service in model_data.values(): + for _class in classes_per_service.values(): + _class.instances.should.equal([]) + + conn = boto3.client("sqs", region_name="us-west-1") + conn.create_queue(QueueName="queue1") + + def test_should_find_bucket(self): + model_data["sqs"]["Queue"].instances.should.have.length_of(1) diff --git a/tests/test_s3/test_s3_file_handles.py b/tests/test_s3/test_s3_file_handles.py index cb4366500..70602145d 100644 --- a/tests/test_s3/test_s3_file_handles.py +++ b/tests/test_s3/test_s3_file_handles.py @@ -168,6 +168,7 @@ class TestS3FileHandleClosures(TestCase): @verify_zero_warnings def test_overwrite_versioned_upload(self): self.s3.put_object("versioned-bucket", "my-key", "x" * 10_000_000) + self.s3.put_object("versioned-bucket", "my-key", "x" * 10_000_000) @verify_zero_warnings def test_multiple_versions_upload(self): @@ -262,6 +263,21 @@ class TestS3FileHandleClosuresUsingMocks(TestCase): version = self.s3.put_object(Bucket="foo", Key="b", Body="s")["VersionId"] self.s3.delete_object(Bucket="foo", Key="b", VersionId=version) + @verify_zero_warnings + def test_update_versioned_object__while_looping(self): + for _ in (1, 2): + with mock_s3(): + self.s3.create_bucket(Bucket="foo") + self.s3.put_bucket_versioning( + Bucket="foo", + VersioningConfiguration={ + "Status": "Enabled", + "MFADelete": "Disabled", + }, + ) + self.s3.put_object(Bucket="foo", Key="bar", Body="stuff") + self.s3.put_object(Bucket="foo", Key="bar", Body="stuff2") + def test_verify_key_can_be_copied_after_disposing(): # https://github.com/getmoto/moto/issues/5588