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)
