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