306 lines
11 KiB
Python
306 lines
11 KiB
Python
import hashlib
|
|
import json
|
|
import re
|
|
|
|
import httplib2
|
|
import pytest
|
|
from googleapiclient.errors import HttpError
|
|
|
|
from conftest import FakeDrive
|
|
from drive_client import (
|
|
DriveClient,
|
|
FolderResolver,
|
|
ResumableUploader,
|
|
classify_http_error,
|
|
escape_query,
|
|
parse_range,
|
|
)
|
|
from errors import AuthError, DriveQuotaError, Interrupted, PermanentError, TransientError, UploadSessionExpired
|
|
|
|
CHUNK = 256 * 1024
|
|
|
|
|
|
class Resp:
|
|
def __init__(self, status, headers=None, body=None):
|
|
self.status_code = status
|
|
self.headers = headers or {}
|
|
self._body = body
|
|
|
|
def json(self):
|
|
if self._body is None:
|
|
raise ValueError("no body")
|
|
return self._body
|
|
|
|
|
|
class FakeUploadServer:
|
|
"""Implements Google's resumable upload protocol in memory."""
|
|
|
|
def __init__(self):
|
|
self.sessions = {}
|
|
self.fail_next_put = [] # exceptions or Resp objects
|
|
self.initiations = 0
|
|
self.put_count = 0
|
|
|
|
def request(self, method, url, data=None, headers=None, timeout=None):
|
|
headers = headers or {}
|
|
if method == "POST":
|
|
self.initiations += 1
|
|
uri = f"https://upload/session/{self.initiations}"
|
|
meta = json.loads(data)
|
|
self.sessions[uri] = {"data": b"", "size": int(headers["X-Upload-Content-Length"]), "meta": meta, "done": None}
|
|
return Resp(200, {"Location": uri})
|
|
session = self.sessions.get(url)
|
|
if session is None:
|
|
return Resp(404)
|
|
content_range = headers["Content-Range"]
|
|
if content_range.startswith("bytes */"):
|
|
return self._status(session)
|
|
self.put_count += 1
|
|
if self.fail_next_put:
|
|
failure = self.fail_next_put.pop(0)
|
|
if isinstance(failure, BaseException):
|
|
# simulate a connection drop after the server stored half the chunk
|
|
session["data"] += data[: len(data) // 2]
|
|
raise failure
|
|
return failure
|
|
start = int(re.match(r"bytes (\d+)-", content_range).group(1))
|
|
assert start == len(session["data"]), "client must continue exactly where the server is"
|
|
session["data"] += data
|
|
return self._status(session)
|
|
|
|
def _status(self, session):
|
|
if len(session["data"]) >= session["size"]:
|
|
if session["done"] is None:
|
|
data = session["data"]
|
|
session["done"] = {
|
|
"id": "file-1",
|
|
"name": session["meta"]["name"],
|
|
"size": str(len(data)),
|
|
"md5Checksum": hashlib.md5(data).hexdigest(),
|
|
"parents": session["meta"]["parents"],
|
|
}
|
|
return Resp(200, body=session["done"])
|
|
headers = {"Range": f"bytes=0-{len(session['data']) - 1}"} if session["data"] else {}
|
|
return Resp(308, headers)
|
|
|
|
|
|
def _upload(server, path, *, session_uri=None, stop_after=None, retries=3):
|
|
sessions, progress = [], []
|
|
calls = {"n": 0}
|
|
|
|
def should_stop():
|
|
calls["n"] += 1
|
|
return stop_after is not None and calls["n"] > stop_after
|
|
|
|
uploader = ResumableUploader(server, chunk_size=CHUNK, retries=retries, sleep=lambda s: None)
|
|
result = uploader.upload(
|
|
path,
|
|
metadata={"name": "Class 01.mp4", "parents": ["F"], "appProperties": {"tg_source": "1:2"}},
|
|
mime_type="video/mp4",
|
|
size=path.stat().st_size,
|
|
session_uri=session_uri,
|
|
on_session=sessions.append,
|
|
on_progress=lambda d, t: progress.append(d),
|
|
should_stop=should_stop,
|
|
)
|
|
return result, sessions, progress
|
|
|
|
|
|
@pytest.fixture
|
|
def big_file(tmp_path):
|
|
path = tmp_path / "video.bin"
|
|
path.write_bytes(bytes(range(256)) * (CHUNK * 3 // 256 + 100)) # 3+ chunks
|
|
return path
|
|
|
|
|
|
def test_full_upload_in_chunks(big_file):
|
|
server = FakeUploadServer()
|
|
result, sessions, progress = _upload(server, big_file)
|
|
assert result["size"] == str(big_file.stat().st_size)
|
|
assert result["md5Checksum"] == hashlib.md5(big_file.read_bytes()).hexdigest()
|
|
assert len(sessions) == 1 and server.put_count == 4
|
|
assert progress[-1] == big_file.stat().st_size
|
|
|
|
|
|
def test_pause_then_resume_same_session(big_file):
|
|
server = FakeUploadServer()
|
|
with pytest.raises(Interrupted):
|
|
_upload(server, big_file, stop_after=2)
|
|
uri = next(iter(server.sessions))
|
|
stored = len(server.sessions[uri]["data"])
|
|
assert 0 < stored < big_file.stat().st_size
|
|
result, sessions, _ = _upload(server, big_file, session_uri=uri)
|
|
assert sessions == [] # no new session was started
|
|
assert server.initiations == 1
|
|
assert result["md5Checksum"] == hashlib.md5(big_file.read_bytes()).hexdigest()
|
|
|
|
|
|
def test_network_error_mid_chunk_resumes_from_server_offset(big_file):
|
|
server = FakeUploadServer()
|
|
server.fail_next_put = [ConnectionError("reset")]
|
|
result, _, _ = _upload(server, big_file)
|
|
assert result["md5Checksum"] == hashlib.md5(big_file.read_bytes()).hexdigest()
|
|
|
|
|
|
def test_server_errors_are_retried_then_raised(big_file):
|
|
server = FakeUploadServer()
|
|
server.fail_next_put = [Resp(503), Resp(500)]
|
|
result, _, _ = _upload(server, big_file)
|
|
assert result["id"] == "file-1"
|
|
server = FakeUploadServer()
|
|
server.fail_next_put = [Resp(503)] * 10
|
|
with pytest.raises(TransientError):
|
|
_upload(server, big_file, retries=2)
|
|
|
|
|
|
def test_expired_session(big_file):
|
|
server = FakeUploadServer()
|
|
# unknown saved session -> a new session is created transparently
|
|
result, sessions, _ = _upload(server, big_file, session_uri="https://upload/session/expired")
|
|
assert len(sessions) == 1 and result["id"] == "file-1"
|
|
server = FakeUploadServer()
|
|
server.fail_next_put = [Resp(404)]
|
|
with pytest.raises(UploadSessionExpired):
|
|
_upload(server, big_file)
|
|
|
|
|
|
def test_quota_and_permission_errors_are_not_retried(big_file):
|
|
server = FakeUploadServer()
|
|
server.fail_next_put = [Resp(403, body={"error": {"errors": [{"reason": "storageQuotaExceeded"}], "message": "full"}})]
|
|
with pytest.raises(DriveQuotaError):
|
|
_upload(server, big_file)
|
|
server = FakeUploadServer()
|
|
server.fail_next_put = [Resp(403, body={"error": {"errors": [{"reason": "insufficientFilePermissions"}]}})]
|
|
with pytest.raises(PermanentError):
|
|
_upload(server, big_file)
|
|
|
|
|
|
def test_chunk_size_must_be_aligned():
|
|
with pytest.raises(ValueError):
|
|
ResumableUploader(None, chunk_size=1000)
|
|
|
|
|
|
def test_helpers():
|
|
assert escape_query("It's a \\ test") == "It\\'s a \\\\ test"
|
|
assert parse_range("bytes=0-1048575") == 1048576
|
|
assert parse_range(None) == 0
|
|
assert isinstance(classify_http_error(401), AuthError)
|
|
assert isinstance(classify_http_error(429), TransientError)
|
|
rate = classify_http_error(403, "userRateLimitExceeded")
|
|
assert isinstance(rate, TransientError) and rate.rate_limited
|
|
assert isinstance(classify_http_error(404), PermanentError)
|
|
assert isinstance(classify_http_error(502), TransientError)
|
|
|
|
|
|
# -- DriveClient retry wrapper ---------------------------------------------------------
|
|
|
|
|
|
class FakeRequest:
|
|
def __init__(self, outcome):
|
|
self.outcome = outcome
|
|
|
|
def execute(self):
|
|
if isinstance(self.outcome, BaseException):
|
|
raise self.outcome
|
|
return self.outcome
|
|
|
|
|
|
def _http_error(status, reason=None):
|
|
body = json.dumps({"error": {"errors": [{"reason": reason}] if reason else [], "message": "x"}}).encode()
|
|
return HttpError(httplib2.Response({"status": status}), body)
|
|
|
|
|
|
def test_execute_retries_transient_then_succeeds():
|
|
outcomes = [_http_error(503), OSError("net down"), {"ok": True}]
|
|
client = DriveClient(None, None, chunk_size=CHUNK, sleep=lambda s: None)
|
|
assert client._execute(lambda: FakeRequest(outcomes.pop(0))) == {"ok": True}
|
|
|
|
|
|
def test_execute_does_not_retry_permission_errors():
|
|
calls = []
|
|
client = DriveClient(None, None, chunk_size=CHUNK, sleep=lambda s: None)
|
|
|
|
def build():
|
|
calls.append(1)
|
|
return FakeRequest(_http_error(403, "insufficientPermissions"))
|
|
|
|
with pytest.raises(PermanentError):
|
|
client._execute(build)
|
|
assert len(calls) == 1
|
|
assert client._execute(lambda: FakeRequest(_http_error(404)), not_found_ok=True) is None
|
|
|
|
|
|
# -- FolderResolver ----------------------------------------------------------------------
|
|
|
|
|
|
class MemoryFolderDB:
|
|
def __init__(self):
|
|
self.folders = {}
|
|
|
|
def get_folder(self, parent, key):
|
|
return self.folders.get((parent, key))
|
|
|
|
def put_folder(self, parent, key, name, folder_id):
|
|
self.folders[(parent, key)] = folder_id
|
|
|
|
def delete_folder(self, parent, key):
|
|
self.folders.pop((parent, key), None)
|
|
|
|
|
|
def _resolver(drive, db=None):
|
|
return FolderResolver(drive, db or MemoryFolderDB(), "root", sleep=lambda s: None)
|
|
|
|
|
|
def test_folder_created_once_and_reused():
|
|
drive = FakeDrive()
|
|
resolver = _resolver(drive)
|
|
first = resolver.resolve(["Misc• Grammar Practice"])
|
|
second = resolver.resolve(["Misc• Grammar Practice"])
|
|
assert first == second and drive.created_folders == 1
|
|
# a fresh process (empty in-memory verification) still reuses the DB mapping
|
|
db = resolver._db
|
|
again = FolderResolver(drive, db, "root").resolve(["Misc• Grammar Practice"])
|
|
assert again == first and drive.created_folders == 1
|
|
|
|
|
|
def test_existing_drive_folder_is_reused_and_duplicates_pick_oldest():
|
|
drive = FakeDrive()
|
|
oldest = drive.add_folder("root", "Error Detection")
|
|
drive.add_folder("root", "Error Detection")
|
|
resolver = _resolver(drive)
|
|
assert resolver.resolve(["Error Detection"]) == oldest.id
|
|
assert drive.created_folders == 0
|
|
|
|
|
|
def test_deleted_cached_folder_is_recreated():
|
|
drive = FakeDrive()
|
|
resolver = _resolver(drive)
|
|
folder_id = resolver.resolve(["Topic"])
|
|
drive.files[folder_id].trashed = True
|
|
fresh = FolderResolver(drive, resolver._db, "root")
|
|
new_id = fresh.resolve(["Topic"])
|
|
assert new_id != folder_id and drive.created_folders == 2
|
|
|
|
|
|
def test_lost_create_response_does_not_duplicate_folder():
|
|
drive = FakeDrive()
|
|
real_create = drive.create_folder
|
|
|
|
def flaky_create(parent, name):
|
|
real_create(parent, name) # Drive created it...
|
|
raise TransientError("...but the response was lost")
|
|
|
|
drive.create_folder = flaky_create
|
|
resolver = _resolver(drive)
|
|
folder_id = resolver.resolve(["Reading Comprehension"])
|
|
assert len(drive.find_folders("root", "Reading Comprehension")) == 1
|
|
assert drive.files[folder_id].name == "Reading Comprehension"
|
|
|
|
|
|
def test_nested_batch_topic_path():
|
|
drive = FakeDrive()
|
|
resolver = _resolver(drive)
|
|
leaf = resolver.resolve(["English ( SSC MAINS )", "Misc• Grammar Practice"])
|
|
batch = drive.find_folders("root", "English ( SSC MAINS )")[0]
|
|
assert drive.files[leaf].parents == (batch.id,)
|