first commit

2026-10-18 23:48:22 +05:30
commit 732a06f65d
38 changed files with 8943 additions and 0 deletions
@@ -0,0 +1,305 @@
+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,)