File: test_writer.py

package info (click to toggle)
python-boto 2.49.0-4.1
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 9,888 kB
  • sloc: python: 86,396; makefile: 112
file content (229 lines) | stat: -rw-r--r-- 8,370 bytes parent folder | download | duplicates (5)
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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
# Copyright (c) 2012 Amazon.com, Inc. or its affiliates.  All Rights Reserved
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish, dis-
# tribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the fol-
# lowing conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#
from hashlib import sha256
import itertools
from boto.compat import StringIO

from tests.unit import unittest
from mock import (
    call,
    Mock,
    sentinel,
)
from nose.tools import assert_equal

from boto.glacier.layer1 import Layer1
from boto.glacier.vault import Vault
from boto.glacier.writer import Writer, resume_file_upload
from boto.glacier.utils import bytes_to_hex, chunk_hashes, tree_hash


def create_mock_vault():
    vault = Mock(spec=Vault)
    vault.layer1 = Mock(spec=Layer1)
    vault.layer1.complete_multipart_upload.return_value = dict(
        ArchiveId=sentinel.archive_id)
    vault.name = sentinel.vault_name
    return vault


def partify(data, part_size):
    for i in itertools.count(0):
        start = i * part_size
        part = data[start:start + part_size]
        if part:
            yield part
        else:
            return


def calculate_mock_vault_calls(data, part_size, chunk_size):
    upload_part_calls = []
    data_tree_hashes = []
    for i, data_part in enumerate(partify(data, part_size)):
        start = i * part_size
        end = start + len(data_part)
        data_part_tree_hash_blob = tree_hash(
            chunk_hashes(data_part, chunk_size))
        data_part_tree_hash = bytes_to_hex(data_part_tree_hash_blob)
        data_part_linear_hash = sha256(data_part).hexdigest()
        upload_part_calls.append(
            call(
                sentinel.vault_name, sentinel.upload_id,
                data_part_linear_hash, data_part_tree_hash,
                (start, end - 1), data_part))
        data_tree_hashes.append(data_part_tree_hash_blob)

    return upload_part_calls, data_tree_hashes


def check_mock_vault_calls(vault, upload_part_calls, data_tree_hashes,
                           data_len):
    vault.layer1.upload_part.assert_has_calls(
        upload_part_calls, any_order=True)
    assert_equal(
        len(upload_part_calls), vault.layer1.upload_part.call_count)

    data_tree_hash = bytes_to_hex(tree_hash(data_tree_hashes))
    vault.layer1.complete_multipart_upload.assert_called_once_with(
        sentinel.vault_name, sentinel.upload_id, data_tree_hash, data_len)


class TestWriter(unittest.TestCase):
    def setUp(self):
        super(TestWriter, self).setUp()
        self.vault = create_mock_vault()
        self.chunk_size = 2 # power of 2
        self.part_size = 4 # power of 2
        upload_id = sentinel.upload_id
        self.writer = Writer(
            self.vault, upload_id, self.part_size, self.chunk_size)

    def check_write(self, write_list):
        for write_data in write_list:
            self.writer.write(write_data)
        self.writer.close()

        data = b''.join(write_list)
        upload_part_calls, data_tree_hashes = calculate_mock_vault_calls(
            data, self.part_size, self.chunk_size)
        check_mock_vault_calls(
            self.vault, upload_part_calls, data_tree_hashes, len(data))

    def test_single_byte_write(self):
        self.check_write([b'1'])

    def test_one_part_write(self):
        self.check_write([b'1234'])

    def test_split_write_1(self):
        self.check_write([b'1', b'234'])

    def test_split_write_2(self):
        self.check_write([b'12', b'34'])

    def test_split_write_3(self):
        self.check_write([b'123', b'4'])

    def test_one_part_plus_one_write(self):
        self.check_write([b'12345'])

    def test_returns_archive_id(self):
        self.writer.write(b'1')
        self.writer.close()
        self.assertEquals(sentinel.archive_id, self.writer.get_archive_id())

    def test_current_tree_hash(self):
        self.writer.write(b'1234')
        self.writer.write(b'567')
        hash_1 = self.writer.current_tree_hash
        self.assertEqual(hash_1,
            b'\x0e\xb0\x11Z\x1d\x1f\n\x10|\xf76\xa6\xf5' +
            b'\x83\xd1\xd5"bU\x0c\x95\xa8<\xf5\x81\xef\x0e\x0f\x95\n\xb7k'
        )

        # This hash will be different, since the content has changed.
        self.writer.write(b'22i3uy')
        hash_2 = self.writer.current_tree_hash
        self.assertEqual(hash_2,
            b'\x7f\xf4\x97\x82U]\x81R\x05#^\xe8\x1c\xd19' +
            b'\xe8\x1f\x9e\xe0\x1aO\xaad\xe5\x06"\xa5\xc0\xa8AdL'
        )
        self.writer.close()

        # Check the final tree hash, post-close.
        final_hash = self.writer.current_tree_hash
        self.assertEqual(final_hash,
            b';\x1a\xb8!=\xf0\x14#\x83\x11\xd5\x0b\x0f' +
            b'\xc7D\xe4\x8e\xd1W\x99z\x14\x06\xb9D\xd0\xf0*\x93\xa2\x8e\xf9'
        )
        # Then assert we don't get a different one on a subsequent call.
        self.assertEqual(final_hash, self.writer.current_tree_hash)

    def test_current_uploaded_size(self):
        self.writer.write(b'1234')
        self.writer.write(b'567')
        size_1 = self.writer.current_uploaded_size
        self.assertEqual(size_1, 4)

        # This hash will be different, since the content has changed.
        self.writer.write(b'22i3uy')
        size_2 = self.writer.current_uploaded_size
        self.assertEqual(size_2, 12)
        self.writer.close()

        # Get the final size, post-close.
        final_size = self.writer.current_uploaded_size
        self.assertEqual(final_size, 13)
        # Then assert we don't get a different one on a subsequent call.
        self.assertEqual(final_size, self.writer.current_uploaded_size)

    def test_upload_id(self):
        self.assertEquals(sentinel.upload_id, self.writer.upload_id)


class TestResume(unittest.TestCase):
    def setUp(self):
        super(TestResume, self).setUp()
        self.vault = create_mock_vault()
        self.chunk_size = 2 # power of 2
        self.part_size = 4 # power of 2

    def check_no_resume(self, data, resume_set=set()):
        fobj = StringIO(data.decode('utf-8'))
        part_hash_map = {}
        for part_index in resume_set:
            start = self.part_size * part_index
            end = start + self.part_size
            part_data = data[start:end]
            part_hash_map[part_index] = tree_hash(
                chunk_hashes(part_data, self.chunk_size))

        resume_file_upload(
            self.vault, sentinel.upload_id, self.part_size, fobj,
            part_hash_map, self.chunk_size)

        upload_part_calls, data_tree_hashes = calculate_mock_vault_calls(
            data, self.part_size, self.chunk_size)
        resume_upload_part_calls = [
            call for part_index, call in enumerate(upload_part_calls)
                    if part_index not in resume_set]
        check_mock_vault_calls(
            self.vault, resume_upload_part_calls, data_tree_hashes, len(data))

    def test_one_part_no_resume(self):
        self.check_no_resume(b'1234')

    def test_two_parts_no_resume(self):
        self.check_no_resume(b'12345678')

    def test_one_part_resume(self):
        self.check_no_resume(b'1234', resume_set=set([0]))

    def test_two_parts_one_resume(self):
        self.check_no_resume(b'12345678', resume_set=set([1]))

    def test_returns_archive_id(self):
        archive_id = resume_file_upload(
            self.vault, sentinel.upload_id, self.part_size, StringIO('1'), {},
            self.chunk_size)
        self.assertEquals(sentinel.archive_id, archive_id)