| OLD | NEW |
| (Empty) |
| 1 # -*- coding: utf-8 -*- | |
| 2 # Copyright 2014 Google Inc. All Rights Reserved. | |
| 3 # | |
| 4 # Licensed under the Apache License, Version 2.0 (the "License"); | |
| 5 # you may not use this file except in compliance with the License. | |
| 6 # You may obtain a copy of the License at | |
| 7 # | |
| 8 # http://www.apache.org/licenses/LICENSE-2.0 | |
| 9 # | |
| 10 # Unless required by applicable law or agreed to in writing, software | |
| 11 # distributed under the License is distributed on an "AS IS" BASIS, | |
| 12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| 13 # See the License for the specific language governing permissions and | |
| 14 # limitations under the License. | |
| 15 """Unit tests for resumable streaming upload functions and classes.""" | |
| 16 | |
| 17 from __future__ import absolute_import | |
| 18 | |
| 19 from hashlib import md5 | |
| 20 import os | |
| 21 import pkgutil | |
| 22 | |
| 23 from gslib.exception import CommandException | |
| 24 from gslib.hashing_helper import CalculateHashesFromContents | |
| 25 from gslib.hashing_helper import CalculateMd5FromContents | |
| 26 from gslib.resumable_streaming_upload import ResumableStreamingJsonUploadWrapper | |
| 27 import gslib.tests.testcase as testcase | |
| 28 from gslib.util import GetJsonResumableChunkSize | |
| 29 from gslib.util import TRANSFER_BUFFER_SIZE | |
| 30 | |
| 31 | |
| 32 _TEST_FILE = 'test.txt' | |
| 33 | |
| 34 | |
| 35 class TestResumableStreamingJsonUploadWrapper(testcase.GsUtilUnitTestCase): | |
| 36 """Unit tests for the TestResumableStreamingJsonUploadWrapper class.""" | |
| 37 | |
| 38 _temp_test_file = None | |
| 39 _temp_test_file_contents = None | |
| 40 _temp_test_file_len = None | |
| 41 | |
| 42 def _GetTestFile(self): | |
| 43 if not self._temp_test_file: | |
| 44 self._temp_test_file_contents = pkgutil.get_data( | |
| 45 'gslib', 'tests/test_data/%s' % _TEST_FILE) | |
| 46 self._temp_test_file = self.CreateTempFile( | |
| 47 file_name=_TEST_FILE, contents=self._temp_test_file_contents) | |
| 48 self._temp_test_file_len = len(self._temp_test_file_contents) | |
| 49 return self._temp_test_file | |
| 50 | |
| 51 def testReadInChunks(self): | |
| 52 tmp_file = self._GetTestFile() | |
| 53 with open(tmp_file, 'rb') as stream: | |
| 54 wrapper = ResumableStreamingJsonUploadWrapper( | |
| 55 stream, TRANSFER_BUFFER_SIZE, test_small_buffer=True) | |
| 56 hash_dict = {'md5': md5()} | |
| 57 # CalculateHashesFromContents reads in chunks, but does not seek. | |
| 58 CalculateHashesFromContents(wrapper, hash_dict) | |
| 59 with open(tmp_file, 'rb') as stream: | |
| 60 actual = CalculateMd5FromContents(stream) | |
| 61 self.assertEqual(actual, hash_dict['md5'].hexdigest()) | |
| 62 | |
| 63 def testReadInChunksWithSeekToBeginning(self): | |
| 64 """Reads one buffer, then seeks to 0 and reads chunks until the end.""" | |
| 65 tmp_file = self._GetTestFile() | |
| 66 for initial_read in (TRANSFER_BUFFER_SIZE - 1, | |
| 67 TRANSFER_BUFFER_SIZE, | |
| 68 TRANSFER_BUFFER_SIZE + 1, | |
| 69 TRANSFER_BUFFER_SIZE * 2 - 1, | |
| 70 TRANSFER_BUFFER_SIZE * 2, | |
| 71 TRANSFER_BUFFER_SIZE * 2 + 1, | |
| 72 TRANSFER_BUFFER_SIZE * 3 - 1, | |
| 73 TRANSFER_BUFFER_SIZE * 3, | |
| 74 TRANSFER_BUFFER_SIZE * 3 + 1): | |
| 75 for buffer_size in (TRANSFER_BUFFER_SIZE - 1, | |
| 76 TRANSFER_BUFFER_SIZE, | |
| 77 TRANSFER_BUFFER_SIZE + 1, | |
| 78 self._temp_test_file_len - 1, | |
| 79 self._temp_test_file_len, | |
| 80 self._temp_test_file_len + 1): | |
| 81 # Can't seek to 0 if the buffer is too small, so we expect an | |
| 82 # exception. | |
| 83 expect_exception = buffer_size < self._temp_test_file_len | |
| 84 with open(tmp_file, 'rb') as stream: | |
| 85 wrapper = ResumableStreamingJsonUploadWrapper( | |
| 86 stream, buffer_size, test_small_buffer=True) | |
| 87 wrapper.read(initial_read) | |
| 88 # CalculateMd5FromContents seeks to 0, reads in chunks, then seeks | |
| 89 # to 0 again. | |
| 90 try: | |
| 91 hex_digest = CalculateMd5FromContents(wrapper) | |
| 92 if expect_exception: | |
| 93 self.fail('Did not get expected CommandException for ' | |
| 94 'initial read size %s, buffer size %s' % | |
| 95 (initial_read, buffer_size)) | |
| 96 except CommandException, e: | |
| 97 if not expect_exception: | |
| 98 self.fail('Got unexpected CommandException "%s" for ' | |
| 99 'initial read size %s, buffer size %s' % | |
| 100 (str(e), initial_read, buffer_size)) | |
| 101 if not expect_exception: | |
| 102 with open(tmp_file, 'rb') as stream: | |
| 103 actual = CalculateMd5FromContents(stream) | |
| 104 self.assertEqual( | |
| 105 actual, hex_digest, | |
| 106 'Digests not equal for initial read size %s, buffer size %s' % | |
| 107 (initial_read, buffer_size)) | |
| 108 | |
| 109 def _testSeekBack(self, initial_reads, buffer_size, seek_back_amount): | |
| 110 """Tests reading then seeking backwards. | |
| 111 | |
| 112 This function simulates an upload that is resumed after a connection break. | |
| 113 It reads one transfer buffer at a time until it reaches initial_position, | |
| 114 then seeks backwards (as if the server did not receive some of the bytes) | |
| 115 and reads to the end of the file, ensuring the data read after the seek | |
| 116 matches the original file. | |
| 117 | |
| 118 Args: | |
| 119 initial_reads: List of integers containing read sizes to perform | |
| 120 before seek. | |
| 121 buffer_size: Maximum buffer size for the wrapper. | |
| 122 seek_back_amount: Number of bytes to seek backward. | |
| 123 | |
| 124 Raises: | |
| 125 AssertionError on wrong data returned by the wrapper. | |
| 126 """ | |
| 127 tmp_file = self._GetTestFile() | |
| 128 initial_position = 0 | |
| 129 for read_size in initial_reads: | |
| 130 initial_position += read_size | |
| 131 self.assertGreaterEqual( | |
| 132 buffer_size, seek_back_amount, | |
| 133 'seek_back_amount must be less than initial position %s ' | |
| 134 '(but was actually: %s)' % (buffer_size, seek_back_amount)) | |
| 135 self.assertLess( | |
| 136 initial_position, self._temp_test_file_len, | |
| 137 'initial_position must be less than test file size %s ' | |
| 138 '(but was actually: %s)' % (self._temp_test_file_len, initial_position)) | |
| 139 | |
| 140 with open(tmp_file, 'rb') as stream: | |
| 141 wrapper = ResumableStreamingJsonUploadWrapper( | |
| 142 stream, buffer_size, test_small_buffer=True) | |
| 143 position = 0 | |
| 144 for read_size in initial_reads: | |
| 145 data = wrapper.read(read_size) | |
| 146 self.assertEqual( | |
| 147 self._temp_test_file_contents[position:position + read_size], | |
| 148 data, 'Data from position %s to %s did not match file contents.' % | |
| 149 (position, position + read_size)) | |
| 150 position += len(data) | |
| 151 wrapper.seek(initial_position - seek_back_amount) | |
| 152 self.assertEqual(wrapper.tell(), | |
| 153 initial_position - seek_back_amount) | |
| 154 data = wrapper.read() | |
| 155 self.assertEqual( | |
| 156 self._temp_test_file_len - (initial_position - seek_back_amount), | |
| 157 len(data), | |
| 158 'Unexpected data length with initial pos %s seek_back_amount %s. ' | |
| 159 'Expected: %s, actual: %s.' % | |
| 160 (initial_position, seek_back_amount, | |
| 161 self._temp_test_file_len - (initial_position - seek_back_amount), | |
| 162 len(data))) | |
| 163 self.assertEqual( | |
| 164 self._temp_test_file_contents[-len(data):], data, | |
| 165 'Data from position %s to EOF did not match file contents.' % | |
| 166 position) | |
| 167 | |
| 168 def testReadSeekAndReadToEOF(self): | |
| 169 """Tests performing reads on the wrapper, seeking, then reading to EOF.""" | |
| 170 for initial_reads in ([1], | |
| 171 [TRANSFER_BUFFER_SIZE - 1], | |
| 172 [TRANSFER_BUFFER_SIZE], | |
| 173 [TRANSFER_BUFFER_SIZE + 1], | |
| 174 [1, TRANSFER_BUFFER_SIZE - 1], | |
| 175 [1, TRANSFER_BUFFER_SIZE], | |
| 176 [1, TRANSFER_BUFFER_SIZE + 1], | |
| 177 [TRANSFER_BUFFER_SIZE - 1, 1], | |
| 178 [TRANSFER_BUFFER_SIZE, 1], | |
| 179 [TRANSFER_BUFFER_SIZE + 1, 1], | |
| 180 [TRANSFER_BUFFER_SIZE - 1, TRANSFER_BUFFER_SIZE - 1], | |
| 181 [TRANSFER_BUFFER_SIZE - 1, TRANSFER_BUFFER_SIZE], | |
| 182 [TRANSFER_BUFFER_SIZE - 1, TRANSFER_BUFFER_SIZE + 1], | |
| 183 [TRANSFER_BUFFER_SIZE, TRANSFER_BUFFER_SIZE - 1], | |
| 184 [TRANSFER_BUFFER_SIZE, TRANSFER_BUFFER_SIZE], | |
| 185 [TRANSFER_BUFFER_SIZE, TRANSFER_BUFFER_SIZE + 1], | |
| 186 [TRANSFER_BUFFER_SIZE + 1, TRANSFER_BUFFER_SIZE - 1], | |
| 187 [TRANSFER_BUFFER_SIZE + 1, TRANSFER_BUFFER_SIZE], | |
| 188 [TRANSFER_BUFFER_SIZE + 1, TRANSFER_BUFFER_SIZE + 1], | |
| 189 [TRANSFER_BUFFER_SIZE, TRANSFER_BUFFER_SIZE, | |
| 190 TRANSFER_BUFFER_SIZE]): | |
| 191 initial_position = 0 | |
| 192 for read_size in initial_reads: | |
| 193 initial_position += read_size | |
| 194 for buffer_size in (initial_position, | |
| 195 initial_position + 1, | |
| 196 initial_position * 2 - 1, | |
| 197 initial_position * 2): | |
| 198 for seek_back_amount in ( | |
| 199 min(TRANSFER_BUFFER_SIZE - 1, initial_position), | |
| 200 min(TRANSFER_BUFFER_SIZE, initial_position), | |
| 201 min(TRANSFER_BUFFER_SIZE + 1, initial_position), | |
| 202 min(TRANSFER_BUFFER_SIZE * 2 - 1, initial_position), | |
| 203 min(TRANSFER_BUFFER_SIZE * 2, initial_position), | |
| 204 min(TRANSFER_BUFFER_SIZE * 2 + 1, initial_position)): | |
| 205 self._testSeekBack(initial_reads, buffer_size, seek_back_amount) | |
| 206 | |
| 207 def testBufferSizeLessThanChunkSize(self): | |
| 208 ResumableStreamingJsonUploadWrapper(None, GetJsonResumableChunkSize()) | |
| 209 try: | |
| 210 ResumableStreamingJsonUploadWrapper(None, GetJsonResumableChunkSize() - 1) | |
| 211 self.fail('Did not get expected CommandException') | |
| 212 except CommandException, e: | |
| 213 self.assertIn('Buffer size must be >= JSON resumable upload', str(e)) | |
| 214 | |
| 215 def testSeekPartialBuffer(self): | |
| 216 """Tests seeking back partially within the buffer.""" | |
| 217 tmp_file = self._GetTestFile() | |
| 218 read_size = TRANSFER_BUFFER_SIZE | |
| 219 with open(tmp_file, 'rb') as stream: | |
| 220 wrapper = ResumableStreamingJsonUploadWrapper( | |
| 221 stream, TRANSFER_BUFFER_SIZE * 3, test_small_buffer=True) | |
| 222 position = 0 | |
| 223 for _ in xrange(3): | |
| 224 data = wrapper.read(read_size) | |
| 225 self.assertEqual( | |
| 226 self._temp_test_file_contents[position:position + read_size], | |
| 227 data, 'Data from position %s to %s did not match file contents.' % | |
| 228 (position, position + read_size)) | |
| 229 position += len(data) | |
| 230 | |
| 231 data = wrapper.read(read_size / 2) | |
| 232 # Buffer contents should now be have contents from: | |
| 233 # read_size/2 through 7*read_size/2. | |
| 234 position = read_size / 2 | |
| 235 wrapper.seek(position) | |
| 236 data = wrapper.read() | |
| 237 self.assertEqual( | |
| 238 self._temp_test_file_contents[-len(data):], data, | |
| 239 'Data from position %s to EOF did not match file contents.' % | |
| 240 position) | |
| 241 | |
| 242 def testSeekEnd(self): | |
| 243 tmp_file = self._GetTestFile() | |
| 244 for buffer_size in (TRANSFER_BUFFER_SIZE - 1, | |
| 245 TRANSFER_BUFFER_SIZE, | |
| 246 TRANSFER_BUFFER_SIZE + 1): | |
| 247 for seek_back in (TRANSFER_BUFFER_SIZE - 1, | |
| 248 TRANSFER_BUFFER_SIZE, | |
| 249 TRANSFER_BUFFER_SIZE + 1): | |
| 250 expect_exception = seek_back > buffer_size | |
| 251 with open(tmp_file, 'rb') as stream: | |
| 252 wrapper = ResumableStreamingJsonUploadWrapper( | |
| 253 stream, buffer_size, test_small_buffer=True) | |
| 254 # Read to the end. | |
| 255 while wrapper.read(TRANSFER_BUFFER_SIZE): | |
| 256 pass | |
| 257 try: | |
| 258 wrapper.seek(seek_back, whence=os.SEEK_END) | |
| 259 if expect_exception: | |
| 260 self.fail('Did not get expected CommandException for ' | |
| 261 'seek_back size %s, buffer size %s' % | |
| 262 (seek_back, buffer_size)) | |
| 263 except CommandException, e: | |
| 264 if not expect_exception: | |
| 265 self.fail('Got unexpected CommandException "%s" for ' | |
| 266 'seek_back size %s, buffer size %s' % | |
| 267 (str(e), seek_back, buffer_size)) | |
| OLD | NEW |