def testRandomFile(self): content = os.urandom(1024) response = responses.Response(responses.POST, "https://foo.bar/quux") response.status = 201 response.headers = { "Location": "https://foo.bar/norf", } responses.add(response) handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.PUT, "https://foo.bar/norf", handler) with temp.AutoTempFilePath() as temp_path: with open(temp_path, mode="wb") as temp_file: temp_file.write(content) flow_id = self._Collect(path=temp_path, signed_url="https://foo.bar/quux") state = flow_test_lib.GetFlowState(self.client_id, flow_id) self.assertNotEmpty(state.encryption_key) encrypted_buf = io.BytesIO(handler.content) decrypted_buf = aead.Decrypt(encrypted_buf, state.encryption_key) self.assertEqual(decrypted_buf.read(), content)
def _testFile(self, content: bytes): # pylint: disable=invalid-name response = responses.Response(responses.POST, "https://foo.bar/quux") response.status = 201 response.headers = { "Location": "https://foo.bar/norf", } responses.add(response) handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.PUT, "https://foo.bar/norf", handler) with temp.AutoTempDirPath(remove_non_empty=True) as tempdir: with open(os.path.join(tempdir, "file"), mode="wb") as file: file.write(content) args = rdf_large_file.CollectLargeFileArgs() args.signed_url = "https://foo.bar/quux" args.path_spec.pathtype = rdf_paths.PathSpec.PathType.OS args.path_spec.path = os.path.join(tempdir, "file") results = list(large_file.CollectLargeFile(args)) self.assertLen(results, 1) self.assertEqual(results[0].session_uri, "https://foo.bar/norf") self.assertEqual(handler.content, content)
def testMultiPartUpload(self): handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.POST, "https://foo.bar/", handler) response = requests.post("https://foo.bar/", data=b"foo", headers={ "Content-Length": "3", "Content-Range": "bytes 0-2/*", }) self.assertEqual(response.status_code, 308) response = requests.post("https://foo.bar/", data=b"bar", headers={ "Content-Length": "3", "Content-Range": "bytes 3-5/*", }) self.assertEqual(response.status_code, 308) response = requests.post("https://foo.bar/", data=b"baz", headers={ "Content-Length": "3", "Content-Range": "bytes 6-8/9", }) self.assertEqual(response.status_code, 201) self.assertEqual(handler.content, b"foobarbaz")
def testSendFileEmpty(self): handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.PUT, "https://foo.bar/qux", handler) session = gcs.UploadSession("https://foo.bar/qux") session.SendFile(io.BytesIO(b"")) self.assertEqual(handler.content, b"")
def testSendFileMultipleChunks(self): handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.PUT, "https://foo.bar/qux", handler) opts = gcs.UploadSession.Opts() opts.chunk_size = 1 session = gcs.UploadSession("https://foo.bar/qux") session.SendFile(io.BytesIO(b"foobar"), opts=opts) self.assertEqual(handler.content, b"foobar")
def testInvalidContentRangeFirstByteOffset(self): handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.POST, "https://foo.bar/", handler) response = requests.post("https://foo.bar/", data=b"foobar", headers={ "Content-Length": "6", "Content-Range": "bytes 3-5/6", }) self.assertEqual(response.status_code, 400) self.assertEqual(response.text, "incorrect first byte offset")
def testInvalidContentRangeSyntax(self): handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.POST, "https://foo.bar/", handler) response = requests.post("https://foo.bar/", data=b"foobar", headers={ "Content-Length": "6", "Content-Range": "bytes ?-?/?", }) self.assertEqual(response.status_code, 400) self.assertEqual(response.text, "invalid content range")
def testIncorrectNumberOfTotalBytes(self): handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.POST, "https://foo.bar/", handler) response = requests.post("https://foo.bar/", data=b"foobar", headers={ "Content-Length": "6", "Content-Range": "bytes 0-5/11", }) self.assertEqual(response.status_code, 400) self.assertEqual(response.text, "incorrect number of total bytes")
def testSinglePartUpload(self): handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.POST, "https://foo.bar/", handler) response = requests.post("https://foo.bar/", data=b"foobar", headers={ "Content-Length": "6", "Content-Range": "bytes 0-5/6", }) self.assertEqual(response.status_code, 201) self.assertEqual(handler.content, b"foobar")
def testSendFileRetrySuccess(self): handler = gcs_test_lib.FakeUploadHandler() responses.add(responses.PUT, "https://foo.bar/qux", status=502) responses.add(responses.PUT, "https://foo.bar/qux", status=503) responses.add(responses.PUT, "https://foo.bar/qux", status=504) responses.add_callback(responses.PUT, "https://foo.bar/qux", handler) opts = gcs.UploadSession.Opts() opts.chunk_size = 1 opts.retry_chunk_attempts = 4 opts.retry_chunk_init_delay = 0.0 session = gcs.UploadSession("https://foo.bar/qux") session.SendFile(io.BytesIO(b"foobar"), opts=opts) self.assertEqual(handler.content, b"foobar")
def testInvalidContentLength(self): handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.POST, "https://foo.bar/", handler) # The `requests` package generally provides its own `Content-Length` header # calculated based on the attached payload. We work around this by rewriting # the value in prepared request. request = requests.Request("POST", "https://foo.bar/", data=b"foobar") request = request.prepare() request.headers["Content-Length"] = "3" request.headers["Content-Range"] = "bytes 0-5/6" with requests.session() as session: response = session.send(request) self.assertEqual(response.status_code, 400) self.assertEqual(response.text, "invalid content length")
def testSendFileRetryFailure(self): handler = gcs_test_lib.FakeUploadHandler() responses.add(responses.PUT, "https://foo.bar/qux", status=502) responses.add(responses.PUT, "https://foo.bar/qux", status=503) responses.add(responses.PUT, "https://foo.bar/qux", status=504) responses.add_callback(responses.PUT, "https://foo.bar/qux", handler) opts = gcs.UploadSession.Opts() opts.chunk_size = 1 opts.retry_chunk_attempts = 3 opts.retry_chunk_init_delay = 0.0 session = gcs.UploadSession("https://foo.bar/qux") with self.assertRaises(gcs.InterruptedResponseError) as context: session.SendFile(io.BytesIO(b"foobar"), opts=opts) self.assertEqual(context.exception.response.status_code, 504)
def testSendFileChunkProgress(self): data = b"foobar" handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.PUT, "https://foo.bar/qux", handler) counter = 0 def Progress() -> None: nonlocal counter counter += 1 opts = gcs.UploadSession.Opts() opts.chunk_size = 1 opts.progress_callback = Progress session = gcs.UploadSession("https://foo.bar/qux") session.SendFile(io.BytesIO(data), opts=opts) self.assertGreaterEqual(counter, len(data))
def testRandomFile(self): content = os.urandom(1024) response = responses.Response(responses.POST, "https://foo.bar/quux") response.status = 201 response.headers = { "Location": "https://foo.bar/norf", } responses.add(response) handler = gcs_test_lib.FakeUploadHandler() responses.add_callback(responses.PUT, "https://foo.bar/norf", handler) with temp.AutoTempFilePath() as temp_path: with open(temp_path, mode="wb") as temp_file: temp_file.write(content) self._Collect(path=temp_path, signed_url="https://foo.bar/quux") self.assertEqual(handler.content, content)