1#!/usr/bin/env python
2# Copyright (c) 2012 Amazon.com, Inc. or its affiliates.  All Rights Reserved
3#
4# Permission is hereby granted, free of charge, to any person obtaining a
5# copy of this software and associated documentation files (the
6# "Software"), to deal in the Software without restriction, including
7# without limitation the rights to use, copy, modify, merge, publish, dis-
8# tribute, sublicense, and/or sell copies of the Software, and to permit
9# persons to whom the Software is furnished to do so, subject to the fol-
10# lowing conditions:
11#
12# The above copyright notice and this permission notice shall be included
13# in all copies or substantial portions of the Software.
14#
15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
16# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
17# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
18# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
19# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
21# IN THE SOFTWARE.
22#
23from boto.compat import StringIO
24from tests.compat import mock, unittest
25
26ANY = mock.ANY
27
28from boto.glacier import vault
29from boto.glacier.job import Job
30from boto.glacier.response import GlacierResponse
31
32
33class TestVault(unittest.TestCase):
34    def setUp(self):
35        self.size_patch = mock.patch('os.path.getsize')
36        self.getsize = self.size_patch.start()
37        self.api = mock.Mock()
38        self.vault = vault.Vault(self.api, None)
39        self.vault.name = 'myvault'
40        self.mock_open = mock.mock_open()
41        stringio = StringIO('content')
42        self.mock_open.return_value.read = stringio.read
43
44    def tearDown(self):
45        self.size_patch.stop()
46
47    @mock.patch('boto.glacier.vault.compute_hashes_from_fileobj',
48                return_value=[b'abc', b'123'])
49    def test_upload_archive_small_file(self, compute_hashes):
50        self.getsize.return_value = 1
51
52        self.api.upload_archive.return_value = {'ArchiveId': 'archive_id'}
53        with mock.patch('boto.glacier.vault.open', self.mock_open,
54                        create=True):
55            archive_id = self.vault.upload_archive(
56                'filename', 'my description')
57        self.assertEqual(archive_id, 'archive_id')
58        self.api.upload_archive.assert_called_with(
59            'myvault', self.mock_open.return_value,
60            mock.ANY, mock.ANY, 'my description')
61
62    def test_small_part_size_is_obeyed(self):
63        self.vault.DefaultPartSize = 2 * 1024 * 1024
64        self.vault.create_archive_writer = mock.Mock()
65
66        self.getsize.return_value = 1
67
68        with mock.patch('boto.glacier.vault.open', self.mock_open,
69                        create=True):
70            self.vault.create_archive_from_file('myfile')
71        # The write should be created with the default part size of the
72        # instance (2 MB).
73        self.vault.create_archive_writer.assert_called_with(
74            description=mock.ANY, part_size=self.vault.DefaultPartSize)
75
76    def test_large_part_size_is_obeyed(self):
77        self.vault.DefaultPartSize = 8 * 1024 * 1024
78        self.vault.create_archive_writer = mock.Mock()
79        self.getsize.return_value = 1
80        with mock.patch('boto.glacier.vault.open', self.mock_open,
81                        create=True):
82            self.vault.create_archive_from_file('myfile')
83        # The write should be created with the default part size of the
84        # instance (8 MB).
85        self.vault.create_archive_writer.assert_called_with(
86            description=mock.ANY, part_size=self.vault.DefaultPartSize)
87
88    def test_part_size_needs_to_be_adjusted(self):
89        # If we have a large file (400 GB)
90        self.getsize.return_value = 400 * 1024 * 1024 * 1024
91        self.vault.create_archive_writer = mock.Mock()
92        # When we try to upload the file.
93        with mock.patch('boto.glacier.vault.open', self.mock_open,
94                        create=True):
95            self.vault.create_archive_from_file('myfile')
96        # We should automatically bump up the part size used to
97        # 64 MB.
98        expected_part_size = 64 * 1024 * 1024
99        self.vault.create_archive_writer.assert_called_with(
100            description=mock.ANY, part_size=expected_part_size)
101
102    def test_retrieve_inventory(self):
103        class FakeResponse(object):
104            status = 202
105
106            def getheader(self, key, default=None):
107                if key == 'x-amz-job-id':
108                    return 'HkF9p6'
109                elif key == 'Content-Type':
110                    return 'application/json'
111
112                return 'something'
113
114            def read(self, amt=None):
115                return b"""{
116  "Action": "ArchiveRetrieval",
117  "ArchiveId": "NkbByEejwEggmBz2fTHgJrg0XBoDfjP4q6iu87-EXAMPLEArchiveId",
118  "ArchiveSizeInBytes": 16777216,
119  "ArchiveSHA256TreeHash": "beb0fe31a1c7ca8c6c04d574ea906e3f97",
120  "Completed": false,
121  "CreationDate": "2012-05-15T17:21:39.339Z",
122  "CompletionDate": "2012-05-15T17:21:43.561Z",
123  "InventorySizeInBytes": null,
124  "JobDescription": "My ArchiveRetrieval Job",
125  "JobId": "HkF9p6",
126  "RetrievalByteRange": "0-16777215",
127  "SHA256TreeHash": "beb0fe31a1c7ca8c6c04d574ea906e3f97b31fd",
128  "SNSTopic": "arn:aws:sns:us-east-1:012345678901:mytopic",
129  "StatusCode": "InProgress",
130  "StatusMessage": "Operation in progress.",
131  "VaultARN": "arn:aws:glacier:us-east-1:012345678901:vaults/examplevault"
132}"""
133
134        raw_resp = FakeResponse()
135        init_resp = GlacierResponse(raw_resp, [('x-amz-job-id', 'JobId')])
136        raw_resp_2 = FakeResponse()
137        desc_resp = GlacierResponse(raw_resp_2, [])
138
139        with mock.patch.object(self.vault.layer1, 'initiate_job',
140                               return_value=init_resp):
141            with mock.patch.object(self.vault.layer1, 'describe_job',
142                                   return_value=desc_resp):
143                # The old/back-compat variant of the call.
144                self.assertEqual(self.vault.retrieve_inventory(), 'HkF9p6')
145
146                # The variant the returns a full ``Job`` object.
147                job = self.vault.retrieve_inventory_job()
148                self.assertTrue(isinstance(job, Job))
149                self.assertEqual(job.id, 'HkF9p6')
150
151
152class TestConcurrentUploads(unittest.TestCase):
153
154    def test_concurrent_upload_file(self):
155        v = vault.Vault(None, None)
156        with mock.patch('boto.glacier.vault.ConcurrentUploader') as c:
157            c.return_value.upload.return_value = 'archive_id'
158            archive_id = v.concurrent_create_archive_from_file(
159                'filename', 'my description')
160            c.return_value.upload.assert_called_with('filename',
161                                                     'my description')
162        self.assertEqual(archive_id, 'archive_id')
163
164    def test_concurrent_upload_forwards_kwargs(self):
165        v = vault.Vault(None, None)
166        with mock.patch('boto.glacier.vault.ConcurrentUploader') as c:
167            c.return_value.upload.return_value = 'archive_id'
168            archive_id = v.concurrent_create_archive_from_file(
169                'filename', 'my description', num_threads=10,
170                part_size=1024 * 1024 * 1024 * 8)
171            c.assert_called_with(None, None, num_threads=10,
172                                 part_size=1024 * 1024 * 1024 * 8)
173
174
175if __name__ == '__main__':
176    unittest.main()
177