322 lines
15 KiB
Python
322 lines
15 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import base64
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
import sqlite3
|
||
|
|
import tempfile
|
||
|
|
import unittest
|
||
|
|
from unittest import mock
|
||
|
|
|
||
|
|
from cmbuyer_client.core.errors import AmbiguousRemoteError, CredentialRemoteError, ManualRemoteError, StateError
|
||
|
|
from cmbuyer_client.core.models import AssetReceipt, ClaimedTask, RenewResult, ScreenshotAsset, SecretToken
|
||
|
|
from cmbuyer_client.localstate.facade import DurableClientGateway
|
||
|
|
from cmbuyer_client.localstate.models import ProfileSettings
|
||
|
|
from cmbuyer_client.localstate.store import LocalStateStore
|
||
|
|
from cmbuyer_client.remote.evidence_sink import HttpEvidenceSink
|
||
|
|
from cmbuyer_client.remote.http_transport import HttpResponse
|
||
|
|
from cmbuyer_client.remote.task_source import HttpTaskSource
|
||
|
|
from tests.core.test_models import ATTEMPT_ID, TASK_ID, claim_wire
|
||
|
|
from tests.localstate.test_store import DEVICE_TOKEN, FakeProtector, PROFILE
|
||
|
|
from tests.remote.test_task_source import DEVICE_ID, FakeTransport, response
|
||
|
|
|
||
|
|
|
||
|
|
PNG = base64.b64decode(
|
||
|
|
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class InspectingSource:
|
||
|
|
def __init__(self, store: LocalStateStore, profile_id: str) -> None:
|
||
|
|
self.store = store
|
||
|
|
self.profile_id = profile_id
|
||
|
|
self.calls = 0
|
||
|
|
self.mode = "success"
|
||
|
|
self.request_ids: list[str] = []
|
||
|
|
|
||
|
|
def claim_next(self, credentials, request):
|
||
|
|
self.calls += 1
|
||
|
|
self.request_ids.append(request.claim_request_id)
|
||
|
|
# HTTP 适配器被调用时，幂等请求必须已经 durable。
|
||
|
|
self.assert_pending(request.claim_request_id)
|
||
|
|
if self.mode == "ambiguous":
|
||
|
|
raise AmbiguousRemoteError("http_result_unknown")
|
||
|
|
if self.mode == "manual":
|
||
|
|
raise ManualRemoteError("claim_requires_manual")
|
||
|
|
return ClaimedTask.from_wire(claim_wire())
|
||
|
|
|
||
|
|
def renew(self, credentials, request):
|
||
|
|
raise AssertionError("not used")
|
||
|
|
|
||
|
|
def assert_pending(self, request_id: str) -> None:
|
||
|
|
snapshot = LocalStateStore(self.store.database_path, FakeProtector()).recovery_snapshot(self.profile_id)
|
||
|
|
if snapshot.pending_claim is None or snapshot.pending_claim.claim_request_id != request_id:
|
||
|
|
raise AssertionError("HTTP happened before durable prepare")
|
||
|
|
|
||
|
|
|
||
|
|
class InspectingSink:
|
||
|
|
def __init__(self, store: LocalStateStore, profile_id: str) -> None:
|
||
|
|
self.store = store
|
||
|
|
self.profile_id = profile_id
|
||
|
|
self.calls = 0
|
||
|
|
|
||
|
|
def upload(self, credentials, upload):
|
||
|
|
self.calls += 1
|
||
|
|
snapshot = LocalStateStore(self.store.database_path, FakeProtector()).recovery_snapshot(self.profile_id)
|
||
|
|
if not snapshot.pending_evidence or snapshot.pending_evidence[0].upload_key != upload.upload_key:
|
||
|
|
raise AssertionError("HTTP happened before durable evidence slot")
|
||
|
|
return AssetReceipt(
|
||
|
|
"63c9f507-7473-4fa6-8d71-8786c34c6301",
|
||
|
|
upload.task_id,
|
||
|
|
upload.attempt_id,
|
||
|
|
upload.kind,
|
||
|
|
upload.privacy_tier,
|
||
|
|
upload.sha256,
|
||
|
|
len(upload.content),
|
||
|
|
"image/png",
|
||
|
|
1,
|
||
|
|
1,
|
||
|
|
upload.captured_at,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class InspectingRenewSource:
|
||
|
|
def __init__(self, store: LocalStateStore, profile_id: str) -> None:
|
||
|
|
self.store = store
|
||
|
|
self.profile_id = profile_id
|
||
|
|
self.calls = 0
|
||
|
|
|
||
|
|
def claim_next(self, credentials, request):
|
||
|
|
raise AssertionError("not used")
|
||
|
|
|
||
|
|
def renew(self, credentials, request):
|
||
|
|
self.calls += 1
|
||
|
|
snapshot = LocalStateStore(self.store.database_path, FakeProtector()).recovery_snapshot(self.profile_id)
|
||
|
|
if snapshot.pending_renew is None or snapshot.pending_renew.renew_request_id != request.renew_request_id:
|
||
|
|
raise AssertionError("HTTP happened before durable renew")
|
||
|
|
return RenewResult(request.task_id, request.attempt_id, request.claim_generation, request.expected_lease_expires_at)
|
||
|
|
|
||
|
|
|
||
|
|
class DurableClientGatewayTests(unittest.TestCase):
|
||
|
|
def setUp(self) -> None:
|
||
|
|
self.directory = tempfile.TemporaryDirectory()
|
||
|
|
self.database = Path(self.directory.name) / "client-state.sqlite3"
|
||
|
|
self.store = LocalStateStore(
|
||
|
|
self.database,
|
||
|
|
FakeProtector(),
|
||
|
|
now=lambda: datetime(2026, 8, 4, 9, 0, tzinfo=timezone.utc),
|
||
|
|
)
|
||
|
|
profile = ProfileSettings(
|
||
|
|
PROFILE,
|
||
|
|
"http://127.0.0.1:8080",
|
||
|
|
DEVICE_ID,
|
||
|
|
"D:/Portable/adb/adb.exe",
|
||
|
|
"192.168.0.173:5555",
|
||
|
|
"wifi",
|
||
|
|
)
|
||
|
|
self.store.save_profile(profile, SecretToken(DEVICE_TOKEN))
|
||
|
|
self.store.start_or_resume_polling(PROFILE)
|
||
|
|
self.source = InspectingSource(self.store, PROFILE)
|
||
|
|
self.sink = InspectingSink(self.store, PROFILE)
|
||
|
|
self.gateway = DurableClientGateway(self.store, self.source, self.sink)
|
||
|
|
|
||
|
|
def tearDown(self) -> None:
|
||
|
|
self.directory.cleanup()
|
||
|
|
|
||
|
|
def test_claim_unknown_replays_same_durable_key_then_commits(self) -> None:
|
||
|
|
self.source.mode = "ambiguous"
|
||
|
|
with self.assertRaises(AmbiguousRemoteError):
|
||
|
|
self.gateway.claim_next(PROFILE)
|
||
|
|
self.source.mode = "success"
|
||
|
|
claimed = self.gateway.claim_next(PROFILE)
|
||
|
|
self.assertEqual(claimed.task.id, TASK_ID)
|
||
|
|
self.assertEqual(self.source.request_ids[0], self.source.request_ids[1])
|
||
|
|
self.assertIsNotNone(self.store.active_claim(PROFILE))
|
||
|
|
|
||
|
|
def test_unknown_claim_2xx_keeps_pending_key_for_real_adapter_replay(self) -> None:
|
||
|
|
transport = FakeTransport(response(201, claim_wire()))
|
||
|
|
gateway = DurableClientGateway(self.store, HttpTaskSource(transport), self.sink)
|
||
|
|
with self.assertRaises(AmbiguousRemoteError):
|
||
|
|
gateway.claim_next(PROFILE)
|
||
|
|
pending = self.store.recovery_snapshot(PROFILE).pending_claim
|
||
|
|
self.assertIsNotNone(pending)
|
||
|
|
transport.response = response(200, claim_wire())
|
||
|
|
claimed = gateway.claim_next(PROFILE)
|
||
|
|
self.assertEqual(claimed.task.id, TASK_ID)
|
||
|
|
sent = [call[3] for call in transport.calls]
|
||
|
|
self.assertEqual(sent[0], sent[1])
|
||
|
|
|
||
|
|
def test_claim_401_allows_token_repair_and_same_key_replay(self) -> None:
|
||
|
|
transport = FakeTransport(HttpResponse(401, (), b""))
|
||
|
|
gateway = DurableClientGateway(self.store, HttpTaskSource(transport), self.sink)
|
||
|
|
with self.assertRaises(CredentialRemoteError):
|
||
|
|
gateway.claim_next(PROFILE)
|
||
|
|
request_id = self.store.recovery_snapshot(PROFILE).pending_claim.claim_request_id
|
||
|
|
self.store.save_profile(self.store.load_profile(PROFILE).settings, SecretToken("c" * 64))
|
||
|
|
transport.response = response(200, claim_wire())
|
||
|
|
gateway.claim_next(PROFILE)
|
||
|
|
self.assertEqual(transport.calls[0][3], transport.calls[1][3])
|
||
|
|
self.assertNotEqual(dict(transport.calls[0][2])["Authorization"], dict(transport.calls[1][2])["Authorization"])
|
||
|
|
self.assertEqual(json.loads(transport.calls[1][3])["claim_request_id"], request_id)
|
||
|
|
|
||
|
|
def test_profile_read_sql_failure_after_prepare_is_fixed_error_and_zero_http(self) -> None:
|
||
|
|
original_connect = self.store._connect
|
||
|
|
calls = 0
|
||
|
|
|
||
|
|
def fail_second_connection():
|
||
|
|
nonlocal calls
|
||
|
|
calls += 1
|
||
|
|
connection = original_connect()
|
||
|
|
if calls == 2:
|
||
|
|
connection.set_authorizer(
|
||
|
|
lambda action, table, *_: sqlite3.SQLITE_DENY
|
||
|
|
if action == sqlite3.SQLITE_READ and table == "profiles"
|
||
|
|
else sqlite3.SQLITE_OK
|
||
|
|
)
|
||
|
|
return connection
|
||
|
|
|
||
|
|
with mock.patch.object(self.store, "_connect", side_effect=fail_second_connection):
|
||
|
|
with self.assertRaisesRegex(StateError, "localstate_read_failed") as captured:
|
||
|
|
self.gateway.claim_next(PROFILE)
|
||
|
|
self.assertEqual(self.source.calls, 0)
|
||
|
|
self.assertNotIn(str(self.database), repr(captured.exception))
|
||
|
|
|
||
|
|
def test_renew_is_durable_before_http(self) -> None:
|
||
|
|
self.gateway.claim_next(PROFILE)
|
||
|
|
source = InspectingRenewSource(self.store, PROFILE)
|
||
|
|
gateway = DurableClientGateway(self.store, source, self.sink)
|
||
|
|
result = gateway.renew(PROFILE)
|
||
|
|
self.assertEqual(result.attempt_id, ATTEMPT_ID)
|
||
|
|
self.assertEqual(source.calls, 1)
|
||
|
|
|
||
|
|
def test_unknown_renew_2xx_keeps_pending_payload_for_replay(self) -> None:
|
||
|
|
self.gateway.claim_next(PROFILE)
|
||
|
|
payload = {
|
||
|
|
"task_id": TASK_ID,
|
||
|
|
"attempt_id": ATTEMPT_ID,
|
||
|
|
"claim_generation": 1,
|
||
|
|
"lease_expires_at": "2026-08-04T09:06:00Z",
|
||
|
|
}
|
||
|
|
transport = FakeTransport(response(201, payload))
|
||
|
|
gateway = DurableClientGateway(self.store, HttpTaskSource(transport), self.sink)
|
||
|
|
with self.assertRaises(AmbiguousRemoteError):
|
||
|
|
gateway.renew(PROFILE)
|
||
|
|
pending = self.store.recovery_snapshot(PROFILE).pending_renew
|
||
|
|
self.assertIsNotNone(pending)
|
||
|
|
transport.response = response(200, payload)
|
||
|
|
gateway.renew(PROFILE)
|
||
|
|
self.assertEqual(transport.calls[0][3], transport.calls[1][3])
|
||
|
|
|
||
|
|
def test_renew_401_allows_bearer_repair_without_changing_claim_payload(self) -> None:
|
||
|
|
self.gateway.claim_next(PROFILE)
|
||
|
|
payload = {
|
||
|
|
"task_id": TASK_ID,
|
||
|
|
"attempt_id": ATTEMPT_ID,
|
||
|
|
"claim_generation": 1,
|
||
|
|
"lease_expires_at": "2026-08-04T09:06:00Z",
|
||
|
|
}
|
||
|
|
transport = FakeTransport(HttpResponse(401, (), b""))
|
||
|
|
gateway = DurableClientGateway(self.store, HttpTaskSource(transport), self.sink)
|
||
|
|
with self.assertRaises(CredentialRemoteError):
|
||
|
|
gateway.renew(PROFILE)
|
||
|
|
self.assertIsNotNone(self.store.recovery_snapshot(PROFILE).pending_renew)
|
||
|
|
self.store.save_profile(self.store.load_profile(PROFILE).settings, SecretToken("c" * 64))
|
||
|
|
transport.response = response(200, payload)
|
||
|
|
gateway.renew(PROFILE)
|
||
|
|
self.assertEqual(transport.calls[0][3], transport.calls[1][3])
|
||
|
|
self.assertNotEqual(dict(transport.calls[0][2])["Authorization"], dict(transport.calls[1][2])["Authorization"])
|
||
|
|
|
||
|
|
def test_unknown_evidence_2xx_keeps_pending_multipart_for_replay(self) -> None:
|
||
|
|
self.gateway.claim_next(PROFILE)
|
||
|
|
path = Path(self.directory.name) / "unknown.png"
|
||
|
|
path.write_bytes(PNG)
|
||
|
|
asset = ScreenshotAsset(path, TASK_ID, ATTEMPT_ID, "2026-08-04T09:01:00Z")
|
||
|
|
transport = FakeTransport(HttpResponse(202, (), b""))
|
||
|
|
gateway = DurableClientGateway(self.store, self.source, HttpEvidenceSink(transport))
|
||
|
|
with self.assertRaises(AmbiguousRemoteError):
|
||
|
|
gateway.upload_evidence(PROFILE, asset)
|
||
|
|
pending = self.store.recovery_snapshot(PROFILE).pending_evidence
|
||
|
|
self.assertEqual(len(pending), 1)
|
||
|
|
digest = hashlib.sha256(PNG).hexdigest()
|
||
|
|
receipt = {
|
||
|
|
"asset_id": "63c9f507-7473-4fa6-8d71-8786c34c6301",
|
||
|
|
"task_id": TASK_ID,
|
||
|
|
"attempt_id": ATTEMPT_ID,
|
||
|
|
"kind": "SKU_PANEL_GATE_1",
|
||
|
|
"privacy_tier": "INTERNAL_RAW",
|
||
|
|
"sha256": digest,
|
||
|
|
"byte_size": len(PNG),
|
||
|
|
"content_type": "image/png",
|
||
|
|
"width_px": 1,
|
||
|
|
"height_px": 1,
|
||
|
|
"captured_at": "2026-08-04T09:01:00Z",
|
||
|
|
}
|
||
|
|
transport.response = HttpResponse(201, (("Content-Type", "application/json"),), json.dumps(receipt).encode())
|
||
|
|
gateway.upload_evidence(PROFILE, asset)
|
||
|
|
self.assertEqual(transport.calls[0][3], transport.calls[1][3])
|
||
|
|
|
||
|
|
def test_evidence_401_allows_bearer_repair_with_same_file_and_multipart(self) -> None:
|
||
|
|
self.gateway.claim_next(PROFILE)
|
||
|
|
path = Path(self.directory.name) / "credential.png"
|
||
|
|
path.write_bytes(PNG)
|
||
|
|
asset = ScreenshotAsset(path, TASK_ID, ATTEMPT_ID, "2026-08-04T09:01:00Z")
|
||
|
|
transport = FakeTransport(HttpResponse(401, (), b""))
|
||
|
|
gateway = DurableClientGateway(self.store, self.source, HttpEvidenceSink(transport))
|
||
|
|
with self.assertRaises(CredentialRemoteError):
|
||
|
|
gateway.upload_evidence(PROFILE, asset)
|
||
|
|
self.assertEqual(len(self.store.recovery_snapshot(PROFILE).pending_evidence), 1)
|
||
|
|
self.store.save_profile(self.store.load_profile(PROFILE).settings, SecretToken("c" * 64))
|
||
|
|
digest = hashlib.sha256(PNG).hexdigest()
|
||
|
|
receipt = {
|
||
|
|
"asset_id": "63c9f507-7473-4fa6-8d71-8786c34c6301",
|
||
|
|
"task_id": TASK_ID,
|
||
|
|
"attempt_id": ATTEMPT_ID,
|
||
|
|
"kind": "SKU_PANEL_GATE_1",
|
||
|
|
"privacy_tier": "INTERNAL_RAW",
|
||
|
|
"sha256": digest,
|
||
|
|
"byte_size": len(PNG),
|
||
|
|
"content_type": "image/png",
|
||
|
|
"width_px": 1,
|
||
|
|
"height_px": 1,
|
||
|
|
"captured_at": "2026-08-04T09:01:00Z",
|
||
|
|
}
|
||
|
|
transport.response = HttpResponse(201, (("Content-Type", "application/json"),), json.dumps(receipt).encode())
|
||
|
|
gateway.upload_evidence(PROFILE, asset)
|
||
|
|
self.assertEqual(transport.calls[0][3], transport.calls[1][3])
|
||
|
|
self.assertNotEqual(dict(transport.calls[0][2])["Authorization"], dict(transport.calls[1][2])["Authorization"])
|
||
|
|
|
||
|
|
def test_equivalent_captured_at_replays_exact_original_multipart_bytes(self) -> None:
|
||
|
|
self.gateway.claim_next(PROFILE)
|
||
|
|
path = Path(self.directory.name) / "exact-replay.png"
|
||
|
|
path.write_bytes(PNG)
|
||
|
|
transport = FakeTransport(HttpResponse(202, (), b""))
|
||
|
|
gateway = DurableClientGateway(self.store, self.source, HttpEvidenceSink(transport))
|
||
|
|
first = ScreenshotAsset(path, TASK_ID, ATTEMPT_ID, "2026-08-04T09:01:00.1Z")
|
||
|
|
equivalent = ScreenshotAsset(path, TASK_ID, ATTEMPT_ID, "2026-08-04T09:01:00.100000Z")
|
||
|
|
with self.assertRaises(AmbiguousRemoteError):
|
||
|
|
gateway.upload_evidence(PROFILE, first)
|
||
|
|
with self.assertRaises(AmbiguousRemoteError):
|
||
|
|
gateway.upload_evidence(PROFILE, equivalent)
|
||
|
|
self.assertEqual(transport.calls[0][3], transport.calls[1][3])
|
||
|
|
|
||
|
|
def test_manual_claim_is_durable_and_never_gets_new_key(self) -> None:
|
||
|
|
self.source.mode = "manual"
|
||
|
|
with self.assertRaises(ManualRemoteError):
|
||
|
|
self.gateway.claim_next(PROFILE)
|
||
|
|
with self.assertRaises(Exception):
|
||
|
|
self.gateway.claim_next(PROFILE)
|
||
|
|
self.assertEqual(self.source.calls, 1)
|
||
|
|
|
||
|
|
def test_evidence_slot_exists_before_http_and_success_never_reuploads(self) -> None:
|
||
|
|
self.gateway.claim_next(PROFILE)
|
||
|
|
path = Path(self.directory.name) / "one.png"
|
||
|
|
path.write_bytes(PNG)
|
||
|
|
asset = ScreenshotAsset(path, TASK_ID, ATTEMPT_ID, "2026-08-04T09:01:00.120000Z")
|
||
|
|
first = self.gateway.upload_evidence(PROFILE, asset)
|
||
|
|
path.write_bytes(PNG + b"changed")
|
||
|
|
second = self.gateway.upload_evidence(PROFILE, asset)
|
||
|
|
self.assertEqual(first, second)
|
||
|
|
self.assertEqual(self.sink.calls, 1)
|