storage/test_integration.py at master · ustudio/storage · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import io
import os
import random
import subprocess
import shutil
import string
import tempfile
import time
import unittest
from urllib.request import urlopen

from storage import get_storage


class IntegrationTests(unittest.TestCase):
    dest_prefix: str
    directory: str

    @classmethod
    def setUpClass(cls) -> None:
        cls.dest_prefix = "".join([random.choice(string.ascii_letters) for i in range(16)])
        cls.directory = tempfile.mkdtemp(prefix="source-root")
        contents = "storage-test-{}".format(time.time()).encode("utf8")

        for i in range(3):
            tempfile.mkstemp(
                prefix="source-empty",
                dir=tempfile.mkdtemp(prefix="source-emptyfiledir", dir=cls.directory))

            handle, _ = tempfile.mkstemp(
                prefix="spaces rock",
                dir=tempfile.mkdtemp(prefix="source-spacedir", dir=cls.directory))
            os.write(handle, contents)
            os.close(handle)

            handle, _ = tempfile.mkstemp(
                prefix="source-contentfile",
                dir=tempfile.mkdtemp(prefix="source-contentdir", dir=cls.directory))
            os.write(handle, contents)
            os.close(handle)

            handle, _ = tempfile.mkstemp(
                prefix="source-mimetyped-file", suffix=".png",
                dir=tempfile.mkdtemp(prefix="source-contentdir", dir=cls.directory))
            os.write(handle, contents)
            os.close(handle)

    @classmethod
    def tearDownClass(cls) -> None:
        shutil.rmtree(cls.directory)

    def assert_transport_handles_directories(self, transport: str) -> None:
        variable = "TEST_STORAGE_{}_URI".format(transport)
        uri = os.getenv(variable, None)

        if not uri:
            raise unittest.SkipTest("Skipping {} - define {} to test".format(transport, variable))

        uri += "/{}".format(self.dest_prefix)
        storage = get_storage(uri)
        print(f"Testing using: {storage.get_sanitized_uri()}")

        print("Transport:", transport)
        print("\t* Uploading")
        storage.load_from_directory(self.directory)

        target_directory = tempfile.mkdtemp(prefix="dest-root")

        print("\t* Downloading")
        storage.save_to_directory(target_directory)

        print("\t* Checking")
        try:
            subprocess.check_output(
                ["diff", "-r", self.directory, target_directory], stderr=subprocess.STDOUT)
        except subprocess.CalledProcessError as error:
            print("Diff output:\n{}".format(error.output))
            raise

    def assert_download_url_generated_correctly(self, transport: str) -> None:
        variable = "TEST_STORAGE_{}_URI".format(transport)
        uri = os.getenv(variable, None)

        if not uri:
            raise unittest.SkipTest("Skipping {} - define {} to test".format(transport, variable))

        uri += "/download-test"
        upload_storage = get_storage(uri)
        print(f"Testing using: {upload_storage.get_sanitized_uri()}")

        print("Transport:", transport)

        upload_storage.load_from_file(io.BytesIO(b"Test data"))

        download_storage = get_storage(uri)

        download_url = download_storage.get_download_url(seconds=3600)

        print("Downloading from download URL")

        with urlopen(download_url) as download:
            assert download.read() == b"Test data"

    def test_file_transport_can_upload_and_download_directories(self) -> None:
        self.assert_transport_handles_directories("FILE")

    def test_ftp_transport_can_upload_and_download_directories(self) -> None:
        self.assert_transport_handles_directories("FTP")

    def test_s3_transport_can_upload_and_download_directories(self) -> None:
        self.assert_transport_handles_directories("S3")

    def test_s3_transport_with_json_credentials_can_upload_and_download_directories(self) -> None:
        self.assert_transport_handles_directories("S3_JSON")

    def test_swift_transport_can_upload_and_download_directories(self) -> None:
        self.assert_transport_handles_directories("SWIFT")

    def test_gs_transport_can_upload_and_download_directories(self) -> None:
        self.assert_transport_handles_directories("GS")

    def test_s3_transport_can_generate_valid_download_urls(self) -> None:
        self.assert_download_url_generated_correctly("S3")

    def test_s3_transport_with_json_credentials_can_generate_valid_download_urls(self) -> None:
        self.assert_download_url_generated_correctly("S3_JSON")

    def test_gs_transport_can_generate_valid_download_urls(self) -> None:
        self.assert_download_url_generated_correctly("GS")