test_uploads: Extract 'use_s3_backend' to test helpers.

zerver/lib/test_helpers.py

@@ -1,11 +1,13 @@
 from contextlib import contextmanager
-from typing import (cast, Any, Callable, Dict, Generator, Iterable, Iterator, List, Mapping,
-                    Optional, Set, Sized, Tuple, Union, IO, Text)
+from typing import (
+    cast, Any, Callable, Dict, Generator, Iterable, Iterator, List, Mapping,
+    Optional, Set, Sized, Tuple, Union, IO, Text, TypeVar
+)
 
 from django.core import signing
 from django.core.urlresolvers import LocaleRegexURLResolver
 from django.conf import settings
-from django.test import TestCase
+from django.test import TestCase, override_settings
 from django.test.client import (
     BOUNDARY, MULTIPART_CONTENT, encode_multipart,
 )
@@ -13,6 +15,8 @@ from django.template import loader
 from django.http import HttpResponse, HttpResponseRedirect
 from django.db.utils import IntegrityError
 
+import zerver.lib.upload
+from zerver.lib.upload import S3UploadBackend, LocalUploadBackend
 from zerver.lib.avatar import avatar_url
 from zerver.lib.cache import get_cache_backend
 from zerver.lib.initial_password import initial_password
@@ -56,6 +60,7 @@ import unittest
 from six.moves import urllib
 from six import binary_type
 from zerver.lib.str_utils import NonBinaryStr
+from moto import mock_s3_deprecated
 
 from contextlib import contextmanager
 import fakeldap
@@ -481,3 +486,18 @@ def load_subdomain_token(response):
     assert isinstance(response, HttpResponseRedirect)
     token = response.url.rsplit('/', 1)[1]
     return signing.loads(token, salt='zerver.views.auth.log_into_subdomain')
+
+FuncT = TypeVar('FuncT', bound=Callable[..., None])
+
+def use_s3_backend(method):
+    # type: (FuncT) -> FuncT
+    @mock_s3_deprecated
+    @override_settings(LOCAL_UPLOADS_DIR=None)
+    def new_method(*args, **kwargs):
+        # type: (*Any, **Any) -> Any
+        zerver.lib.upload.upload_backend = S3UploadBackend()
+        try:
+            return method(*args, **kwargs)
+        finally:
+            zerver.lib.upload.upload_backend = LocalUploadBackend()
+    return new_method