tools/telemetry/third_party/gsutilz/gslib/tests/test_resumable_streaming.py

   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))