1251 lines
50 KiB
Python
1251 lines
50 KiB
Python
import base64
|
||
import json
|
||
import logging
|
||
import os
|
||
import threading
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
from unittest.mock import patch
|
||
|
||
import auto_pool_maintainer as apm
|
||
|
||
|
||
class DummyResponse:
|
||
def __init__(self, status_code: int, *, text: str = "", payload=None):
|
||
self.status_code = status_code
|
||
self.text = text
|
||
self._payload = payload if payload is not None else {}
|
||
self.headers = {}
|
||
self.url = "https://auth.openai.com/email-verification"
|
||
|
||
def json(self):
|
||
if isinstance(self._payload, Exception):
|
||
raise self._payload
|
||
return self._payload
|
||
|
||
|
||
def build_test_jwt(payload: dict) -> str:
|
||
header = base64.urlsafe_b64encode(json.dumps({"alg": "none", "typ": "JWT"}).encode("utf-8")).rstrip(b"=").decode("ascii")
|
||
body = base64.urlsafe_b64encode(json.dumps(payload).encode("utf-8")).rstrip(b"=").decode("ascii")
|
||
return f"{header}.{body}.signature"
|
||
|
||
|
||
class FlowHelperTests(unittest.TestCase):
|
||
def test_request_with_local_retry_writes_flow_trace_log_with_redaction(self):
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
trace_path = os.path.join(tmp_dir, "flow-trace.jsonl")
|
||
recorder = apm.FlowTraceRecorder(trace_path, reveal_sensitive=False, body_limit=512)
|
||
|
||
class FakeSession:
|
||
def __init__(self):
|
||
self.cookies = []
|
||
|
||
def post(self, url, **kwargs):
|
||
response = DummyResponse(429, text='{"error":"rate_limit"}')
|
||
response.url = url
|
||
response.headers = {
|
||
"content-type": "application/json",
|
||
"set-cookie": "session=secret-cookie",
|
||
}
|
||
return response
|
||
|
||
response, reason = apm.request_with_local_retry(
|
||
FakeSession(), # type: ignore[arg-type]
|
||
"post",
|
||
"https://auth.openai.com/api/accounts/authorize/continue",
|
||
retry_attempts=1,
|
||
error_prefix="authorize_continue_request",
|
||
flow_trace=recorder,
|
||
headers={
|
||
"Authorization": "Bearer super-secret-token",
|
||
"Cookie": "session=plain-cookie",
|
||
"x-test": "ok",
|
||
},
|
||
json={"password": "PlainPassword123", "username": "trace@example.com"},
|
||
timeout=30,
|
||
verify=False,
|
||
)
|
||
|
||
self.assertEqual(reason, "")
|
||
self.assertIsNotNone(response)
|
||
|
||
with open(trace_path, "r", encoding="utf-8") as trace_file:
|
||
events = [json.loads(line) for line in trace_file if line.strip()]
|
||
|
||
event_names = [event["event"] for event in events]
|
||
self.assertIn("http_attempt", event_names)
|
||
self.assertIn("http_response", event_names)
|
||
|
||
attempt_event = next(event for event in events if event["event"] == "http_attempt")
|
||
response_event = next(event for event in events if event["event"] == "http_response")
|
||
|
||
self.assertEqual(attempt_event["request"]["url"], "https://auth.openai.com/api/accounts/authorize/continue")
|
||
self.assertEqual(response_event["response"]["status_code"], 429)
|
||
serialized = json.dumps(events, ensure_ascii=False)
|
||
self.assertNotIn("PlainPassword123", serialized)
|
||
self.assertNotIn("super-secret-token", serialized)
|
||
self.assertNotIn("plain-cookie", serialized)
|
||
|
||
def test_build_chatgpt_session_token_result_uses_callback_code(self):
|
||
access_token = build_test_jwt(
|
||
{
|
||
"email": "jwt@example.com",
|
||
"exp": 1760000000,
|
||
"https://api.openai.com/auth": {"chatgpt_account_id": "acct_123"},
|
||
}
|
||
)
|
||
|
||
class FakeSession:
|
||
def __init__(self):
|
||
self.calls = []
|
||
|
||
def get(self, url, **kwargs):
|
||
self.calls.append((url, kwargs))
|
||
if url == "https://chatgpt.com/api/auth/callback/openai?code=oauth-code":
|
||
return DummyResponse(200)
|
||
if url == "https://chatgpt.com/api/auth/session":
|
||
return DummyResponse(200, payload={"accessToken": access_token, "user": {"email": "jwt@example.com"}})
|
||
raise AssertionError(f"unexpected url: {url}")
|
||
|
||
session = FakeSession()
|
||
|
||
result = apm.build_chatgpt_session_token_result(
|
||
session=session, # type: ignore[arg-type]
|
||
auth_code="oauth-code",
|
||
chatgpt_base="https://chatgpt.com",
|
||
)
|
||
|
||
self.assertIsNotNone(result)
|
||
self.assertEqual(result["access_token"], access_token)
|
||
self.assertEqual(result["email"], "jwt@example.com")
|
||
self.assertEqual(result["account_id"], "acct_123")
|
||
self.assertEqual(result["exp"], 1760000000)
|
||
self.assertEqual(
|
||
[call[0] for call in session.calls],
|
||
[
|
||
"https://chatgpt.com/api/auth/callback/openai?code=oauth-code",
|
||
"https://chatgpt.com/api/auth/session",
|
||
],
|
||
)
|
||
|
||
def test_build_chatgpt_session_token_result_preserves_callback_query_params(self):
|
||
access_token = build_test_jwt(
|
||
{
|
||
"email": "jwt@example.com",
|
||
"exp": 1760000000,
|
||
"https://api.openai.com/auth": {"chatgpt_account_id": "acct_123"},
|
||
}
|
||
)
|
||
|
||
class FakeSession:
|
||
def __init__(self):
|
||
self.calls = []
|
||
|
||
def get(self, url, **kwargs):
|
||
self.calls.append((url, kwargs))
|
||
if (
|
||
url
|
||
== "https://chatgpt.com/api/auth/callback/openai"
|
||
"?code=oauth-code&scope=openid+email+profile+offline_access&state=oauth-state"
|
||
):
|
||
return DummyResponse(200)
|
||
if url == "https://chatgpt.com/api/auth/session":
|
||
return DummyResponse(200, payload={"accessToken": access_token, "user": {"email": "jwt@example.com"}})
|
||
raise AssertionError(f"unexpected url: {url}")
|
||
|
||
session = FakeSession()
|
||
|
||
result = apm.build_chatgpt_session_token_result(
|
||
session=session, # type: ignore[arg-type]
|
||
auth_code="oauth-code",
|
||
callback_params={
|
||
"code": "oauth-code",
|
||
"scope": "openid email profile offline_access",
|
||
"state": "oauth-state",
|
||
},
|
||
chatgpt_base="https://chatgpt.com",
|
||
)
|
||
|
||
self.assertIsNotNone(result)
|
||
self.assertEqual(result["access_token"], access_token)
|
||
self.assertEqual(
|
||
[call[0] for call in session.calls],
|
||
[
|
||
"https://chatgpt.com/api/auth/callback/openai?code=oauth-code&scope=openid+email+profile+offline_access&state=oauth-state",
|
||
"https://chatgpt.com/api/auth/session",
|
||
],
|
||
)
|
||
|
||
def test_build_chatgpt_session_token_result_finds_nested_jwt(self):
|
||
access_token = build_test_jwt(
|
||
{
|
||
"email": "jwt@example.com",
|
||
"exp": 1760000000,
|
||
"https://api.openai.com/auth": {"chatgpt_account_id": "acct_123"},
|
||
}
|
||
)
|
||
|
||
class FakeSession:
|
||
def __init__(self):
|
||
self.calls = []
|
||
|
||
def get(self, url, **kwargs):
|
||
self.calls.append((url, kwargs))
|
||
if url == "https://chatgpt.com/api/auth/session":
|
||
return DummyResponse(
|
||
200,
|
||
payload={
|
||
"user": {"email": "jwt@example.com"},
|
||
"session": {"tokens": [{"kind": "bearer", "value": access_token}]},
|
||
},
|
||
)
|
||
raise AssertionError(f"unexpected url: {url}")
|
||
|
||
session = FakeSession()
|
||
|
||
result = apm.build_chatgpt_session_token_result(
|
||
session=session, # type: ignore[arg-type]
|
||
auth_code="",
|
||
chatgpt_base="https://chatgpt.com",
|
||
)
|
||
|
||
self.assertIsNotNone(result)
|
||
self.assertEqual(result["access_token"], access_token)
|
||
self.assertEqual(result["email"], "jwt@example.com")
|
||
|
||
def test_is_transient_flow_error(self):
|
||
self.assertTrue(apm.is_transient_flow_error("oauth_step_http_503"))
|
||
self.assertTrue(apm.is_transient_flow_error("authorize_exception:timed out"))
|
||
self.assertFalse(apm.is_transient_flow_error("email_otp_validate_http_400"))
|
||
|
||
def test_parse_otp_validate_order(self):
|
||
self.assertEqual(apm.parse_otp_validate_order("normal,sentinel"), ("normal", "sentinel"))
|
||
self.assertEqual(apm.parse_otp_validate_order("sentinel,normal"), ("sentinel", "normal"))
|
||
self.assertEqual(apm.parse_otp_validate_order("invalid"), ("normal", "sentinel"))
|
||
|
||
def test_requires_phone_verification(self):
|
||
payload = {
|
||
"page": {"type": "phone_verification"},
|
||
"continue_url": "/add-phone",
|
||
}
|
||
self.assertTrue(apm.requires_phone_verification(payload, ""))
|
||
self.assertFalse(apm.requires_phone_verification({"page": {"type": "email_otp_verification"}}, ""))
|
||
|
||
def test_resolve_loop_interval_seconds(self):
|
||
self.assertEqual(apm.resolve_loop_interval_seconds({}, None), 60.0)
|
||
self.assertEqual(apm.resolve_loop_interval_seconds({"maintainer": {"loop_interval_seconds": 12}}, None), 12.0)
|
||
self.assertEqual(apm.resolve_loop_interval_seconds({"maintainer": {"loop_interval_seconds": 1}}, None), 5.0)
|
||
self.assertEqual(apm.resolve_loop_interval_seconds({}, 8.5), 8.5)
|
||
|
||
def test_parse_loop_next_check_in_seconds_from_log_line(self):
|
||
line = "2026-03-27 21:33:42 | INFO | 循环模式休眠 60.0s 后再次检查号池"
|
||
with patch("api_server.time.time", return_value=apm.dt.datetime(2026, 3, 27, 21, 34, 0).timestamp()):
|
||
import api_server as aps
|
||
|
||
remain = aps.parse_loop_next_check_in_seconds([line])
|
||
self.assertEqual(remain, 42)
|
||
|
||
def test_api_server_run_state_read_write_and_clear(self):
|
||
import api_server as aps
|
||
|
||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||
fake_state = aps.Path(tmp_dir) / "run_state.json"
|
||
with patch.object(aps, "RUN_STATE_FILE", fake_state):
|
||
aps.save_run_state(12345, "loop")
|
||
state = aps.load_run_state()
|
||
self.assertEqual(state.get("pid"), 12345)
|
||
self.assertEqual(state.get("mode"), "loop")
|
||
aps.clear_run_state()
|
||
self.assertFalse(fake_state.exists())
|
||
|
||
def test_api_server_is_pid_running_current_process(self):
|
||
import api_server as aps
|
||
|
||
self.assertTrue(aps.is_pid_running(os.getpid()))
|
||
self.assertFalse(aps.is_pid_running(99999999))
|
||
|
||
def test_analyze_usage_status_marks_quota_and_threshold(self):
|
||
body = {
|
||
"rate_limit": {
|
||
"allowed": True,
|
||
"limit_reached": False,
|
||
"primary_window": {"used_percent": 85},
|
||
"secondary_window": {"used_percent": 99},
|
||
}
|
||
}
|
||
usage = apm.analyze_usage_status(status_code=200, body_obj=body, body_text="", used_percent_threshold=80)
|
||
self.assertEqual(usage["used_percent"], 99.0)
|
||
self.assertTrue(usage["over_threshold"])
|
||
self.assertTrue(usage["is_quota"])
|
||
self.assertFalse(usage["is_healthy"])
|
||
|
||
def test_analyze_usage_status_marks_healthy(self):
|
||
body = {
|
||
"rate_limit": {
|
||
"allowed": True,
|
||
"limit_reached": False,
|
||
"primary_window": {"used_percent": 35},
|
||
}
|
||
}
|
||
usage = apm.analyze_usage_status(status_code=200, body_obj=body, body_text="", used_percent_threshold=80)
|
||
self.assertEqual(usage["used_percent"], 35.0)
|
||
self.assertFalse(usage["over_threshold"])
|
||
self.assertFalse(usage["is_quota"])
|
||
self.assertTrue(usage["is_healthy"])
|
||
|
||
def test_decide_clean_action(self):
|
||
self.assertEqual(apm.decide_clean_action(status_code=401, disabled=False, is_quota=False, over_threshold=False), "delete")
|
||
self.assertEqual(apm.decide_clean_action(status_code=200, disabled=False, is_quota=True, over_threshold=False), "disable")
|
||
self.assertEqual(apm.decide_clean_action(status_code=200, disabled=True, is_quota=False, over_threshold=False), "enable")
|
||
self.assertEqual(apm.decide_clean_action(status_code=None, disabled=False, is_quota=False, over_threshold=False), "keep")
|
||
|
||
def test_get_candidates_count_excludes_disabled_items(self):
|
||
files = [
|
||
{"type": "codex", "disabled": False},
|
||
{"type": "codex", "disabled": True},
|
||
{"type": "codex", "disabled": "false"},
|
||
{"type": "codex", "status": "disabled"},
|
||
{"type": "claude", "disabled": False},
|
||
]
|
||
total, candidates = apm.get_candidates_count_from_files(files, "codex")
|
||
self.assertEqual(total, 5)
|
||
self.assertEqual(candidates, 2)
|
||
|
||
def test_select_probe_candidates_returns_all_when_sample_size_disabled(self):
|
||
candidates = [{"name": "a"}, {"name": "b"}, {"name": "c"}]
|
||
selected = apm.select_probe_candidates(candidates, sample_size=0, rng=apm.random.Random(1))
|
||
self.assertEqual([item["name"] for item in selected], ["a", "b", "c"])
|
||
|
||
def test_select_probe_candidates_returns_random_subset(self):
|
||
candidates = [{"name": "a"}, {"name": "b"}, {"name": "c"}, {"name": "d"}, {"name": "e"}]
|
||
selected = apm.select_probe_candidates(candidates, sample_size=2, rng=apm.random.Random(7))
|
||
self.assertEqual([item["name"] for item in selected], ["c", "b"])
|
||
|
||
def test_run_clean_401_passes_sample_size_to_async_cleanup(self):
|
||
conf = {
|
||
"clean": {
|
||
"base_url": "https://example.test",
|
||
"token": "pw",
|
||
"sample_size": 3,
|
||
}
|
||
}
|
||
captured = {}
|
||
|
||
async def fake_run_clean_401_async(**kwargs):
|
||
captured.update(kwargs)
|
||
return {"action_total": 0}
|
||
|
||
with patch.object(apm, "aiohttp", object()), patch.object(apm, "run_clean_401_async", fake_run_clean_401_async):
|
||
result = apm.run_clean_401(conf, logging.getLogger("test-clean-sample"))
|
||
|
||
self.assertEqual(captured["sample_size"], 3)
|
||
self.assertEqual(result["action_total"], 0)
|
||
|
||
def test_mail_provider_session_reuses_same_thread_and_isolates_cross_thread(self):
|
||
provider = apm.SelfHostedMailApiProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-mail-session"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain="x.test",
|
||
)
|
||
main_session_first = provider._session()
|
||
main_session_second = provider._session()
|
||
self.assertIs(main_session_first, main_session_second)
|
||
|
||
holder = {}
|
||
|
||
def worker() -> None:
|
||
holder["thread_session_first"] = provider._session()
|
||
holder["thread_session_second"] = provider._session()
|
||
|
||
t = threading.Thread(target=worker)
|
||
t.start()
|
||
t.join(timeout=3)
|
||
self.assertIn("thread_session_first", holder)
|
||
self.assertIs(holder["thread_session_first"], holder["thread_session_second"])
|
||
self.assertIsNot(main_session_first, holder["thread_session_first"])
|
||
|
||
def test_self_hosted_mail_domain_normalization_removes_leading_dot(self):
|
||
provider = apm.SelfHostedMailApiProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-self-hosted-domain"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain=".qzz.io",
|
||
)
|
||
mailbox = provider.create_mailbox()
|
||
self.assertIsNotNone(mailbox)
|
||
self.assertEqual(provider.domain, "qzz.io")
|
||
self.assertNotIn("@.", mailbox.email if mailbox else "")
|
||
|
||
def test_yyds_mail_domain_normalization_removes_leading_dot(self):
|
||
provider = apm.YYDSMailProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-yyds-domain"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain=".qzz.io",
|
||
)
|
||
self.assertEqual(provider.domain, "qzz.io")
|
||
|
||
def test_self_hosted_provider_accepts_code_without_openai_keywords(self):
|
||
provider = apm.SelfHostedMailApiProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-self-hosted-code"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain="qzz.io",
|
||
)
|
||
provider._fetch_latest_email = lambda _email: { # type: ignore[method-assign]
|
||
"subject": "您的登录验证码",
|
||
"text": "验证码：123456，请在页面输入",
|
||
}
|
||
codes = provider.poll_verification_codes(
|
||
apm.Mailbox(email="u@qzz.io"),
|
||
seen_ids=set(),
|
||
)
|
||
self.assertEqual(codes, ["123456"])
|
||
|
||
def test_self_hosted_provider_logs_non_200_fetch_response(self):
|
||
logger_name = "test-self-hosted-fetch-warning"
|
||
provider = apm.SelfHostedMailApiProvider(
|
||
proxy="",
|
||
logger=logging.getLogger(logger_name),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain="qzz.io",
|
||
)
|
||
|
||
class FakeResponse:
|
||
status_code = 401
|
||
text = "无效的邮箱地址凭据"
|
||
|
||
def json(self):
|
||
return {}
|
||
|
||
class FakeSession:
|
||
@staticmethod
|
||
def get(*args, **kwargs):
|
||
return FakeResponse()
|
||
|
||
provider._thread_local.session = FakeSession()
|
||
with self.assertLogs(logger_name, level="WARNING") as captured:
|
||
mail_obj = provider._fetch_latest_email("u@qzz.io")
|
||
|
||
self.assertIsNone(mail_obj)
|
||
self.assertTrue(any("401" in line and "无效的邮箱地址凭据" in line for line in captured.output))
|
||
|
||
def test_yyds_provider_accepts_code_without_openai_keywords(self):
|
||
provider = apm.YYDSMailProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-yyds-code"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain="qzz.io",
|
||
)
|
||
provider._fetch_messages = lambda _token: [{"id": "m-1"}] # type: ignore[method-assign]
|
||
provider._fetch_message_detail = lambda _token, _mid: { # type: ignore[method-assign]
|
||
"subject": "邮箱验证码",
|
||
"text": "本次验证码 654321，5 分钟内有效",
|
||
}
|
||
codes = provider.poll_verification_codes(
|
||
apm.Mailbox(email="u@qzz.io", token="tkn"),
|
||
seen_ids=set(),
|
||
)
|
||
self.assertEqual(codes, ["654321"])
|
||
|
||
def test_yyds_provider_accepts_code_from_inline_message_without_detail(self):
|
||
provider = apm.YYDSMailProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-yyds-inline-code"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain="qzz.io",
|
||
)
|
||
provider._fetch_messages = lambda _token: [ # type: ignore[method-assign]
|
||
{"id": "m-1", "subject": "邮箱验证码", "intro": "本次验证码 112233，5 分钟内有效"}
|
||
]
|
||
provider._fetch_message_detail = lambda _token, _mid: None # type: ignore[method-assign]
|
||
codes = provider.poll_verification_codes(
|
||
apm.Mailbox(email="u@qzz.io", token="tkn"),
|
||
seen_ids=set(),
|
||
)
|
||
self.assertEqual(codes, ["112233"])
|
||
|
||
def test_yyds_provider_normalizes_prefixed_message_id_for_detail_fetch(self):
|
||
provider = apm.YYDSMailProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-yyds-message-id"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain="qzz.io",
|
||
)
|
||
provider._fetch_messages = lambda _token: [{"id": "/messages/m-1"}] # type: ignore[method-assign]
|
||
detail_call = {}
|
||
|
||
def fake_fetch_detail(_token, message_id):
|
||
detail_call["message_id"] = message_id
|
||
return {
|
||
"subject": "邮箱验证码",
|
||
"text": "本次验证码 445566，5 分钟内有效",
|
||
}
|
||
|
||
provider._fetch_message_detail = fake_fetch_detail # type: ignore[method-assign]
|
||
codes = provider.poll_verification_codes(
|
||
apm.Mailbox(email="u@qzz.io", token="tkn"),
|
||
seen_ids=set(),
|
||
)
|
||
self.assertEqual(codes, ["445566"])
|
||
self.assertEqual(detail_call.get("message_id"), "m-1")
|
||
|
||
def test_yyds_provider_fetch_messages_reads_nested_messages_array(self):
|
||
provider = apm.YYDSMailProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-yyds-nested-messages"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain="qzz.io",
|
||
)
|
||
|
||
class FakeResponse:
|
||
status_code = 200
|
||
content = b"1"
|
||
|
||
@staticmethod
|
||
def json():
|
||
return {
|
||
"success": True,
|
||
"data": {
|
||
"messages": [
|
||
{"id": "m-1", "subject": "邮箱验证码", "createdAt": "2026-03-28T16:00:00Z"}
|
||
]
|
||
},
|
||
}
|
||
|
||
class FakeSession:
|
||
@staticmethod
|
||
def get(*args, **kwargs):
|
||
return FakeResponse()
|
||
|
||
provider._thread_local.session = FakeSession()
|
||
messages = provider._fetch_messages("tkn")
|
||
self.assertEqual(messages, [{"id": "m-1", "subject": "邮箱验证码", "createdAt": "2026-03-28T16:00:00Z"}])
|
||
|
||
def test_self_hosted_provider_prefers_domains_over_domain(self):
|
||
provider = apm.SelfHostedMailApiProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-self-hosted-domains-priority"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain="fallback.test",
|
||
domains=["a.test", "b.test"],
|
||
failure_threshold=2,
|
||
failure_cooldown_seconds=30.0,
|
||
)
|
||
|
||
self.assertEqual(provider.domains, ["a.test", "b.test"])
|
||
mailbox = provider.create_mailbox()
|
||
self.assertIsNotNone(mailbox)
|
||
self.assertTrue((mailbox.email if mailbox else "").endswith("@a.test"))
|
||
self.assertEqual(mailbox.domain if mailbox else "", "a.test")
|
||
|
||
def test_self_hosted_provider_rotates_domains_in_order(self):
|
||
provider = apm.SelfHostedMailApiProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-self-hosted-rotate"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain="fallback.test",
|
||
domains=["a.test", "b.test", "c.test"],
|
||
failure_threshold=2,
|
||
failure_cooldown_seconds=30.0,
|
||
)
|
||
|
||
first = provider.create_mailbox()
|
||
second = provider.create_mailbox()
|
||
third = provider.create_mailbox()
|
||
|
||
self.assertEqual([first.domain, second.domain, third.domain], ["a.test", "b.test", "c.test"])
|
||
|
||
def test_self_hosted_provider_skips_domain_in_cooldown(self):
|
||
provider = apm.SelfHostedMailApiProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-self-hosted-cooldown"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain="fallback.test",
|
||
domains=["a.test", "b.test"],
|
||
failure_threshold=2,
|
||
failure_cooldown_seconds=60.0,
|
||
)
|
||
|
||
provider.note_domain_failure("a.test", stage="create_mailbox")
|
||
provider.note_domain_failure("a.test", stage="create_mailbox")
|
||
|
||
mailbox = provider.create_mailbox()
|
||
self.assertIsNotNone(mailbox)
|
||
self.assertEqual(mailbox.domain if mailbox else "", "b.test")
|
||
|
||
def test_self_hosted_provider_reuses_domain_after_cooldown_expires(self):
|
||
provider = apm.SelfHostedMailApiProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-self-hosted-cooldown-expire"),
|
||
api_base="https://example.test",
|
||
api_key="k",
|
||
domain="fallback.test",
|
||
domains=["a.test", "b.test"],
|
||
failure_threshold=1,
|
||
failure_cooldown_seconds=5.0,
|
||
)
|
||
|
||
provider.note_domain_failure("a.test", stage="create_mailbox")
|
||
provider.domain_cooldown_until["a.test"] = time.time() - 1
|
||
|
||
mailbox = provider.create_mailbox()
|
||
self.assertIsNotNone(mailbox)
|
||
self.assertEqual(mailbox.domain if mailbox else "", "a.test")
|
||
|
||
def test_cfmail_provider_create_mailbox_uses_next_available_domain(self):
|
||
provider = apm.CfmailProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-cfmail-provider"),
|
||
api_base="https://mail.example.com",
|
||
api_key="pw",
|
||
domain="",
|
||
domains=["a.test", "b.test"],
|
||
failure_threshold=2,
|
||
failure_cooldown_seconds=60.0,
|
||
)
|
||
|
||
provider._create_address_for_domain = lambda domain: apm.Mailbox( # type: ignore[method-assign]
|
||
email=f"oc123@{domain}",
|
||
token="jwt",
|
||
domain=domain,
|
||
failure_target=domain,
|
||
)
|
||
|
||
first = provider.create_mailbox()
|
||
second = provider.create_mailbox()
|
||
|
||
self.assertIsNotNone(first)
|
||
self.assertIsNotNone(second)
|
||
self.assertEqual((first.domain, second.domain), ("a.test", "b.test"))
|
||
|
||
def test_cfmail_provider_extracts_code_from_raw_and_metadata(self):
|
||
provider = apm.CfmailProvider(
|
||
proxy="",
|
||
logger=logging.getLogger("test-cfmail-code"),
|
||
api_base="https://mail.example.com",
|
||
api_key="pw",
|
||
domain="",
|
||
domains=["a.test"],
|
||
failure_threshold=2,
|
||
failure_cooldown_seconds=60.0,
|
||
)
|
||
provider._fetch_cfmail_messages = lambda _mailbox: [ # type: ignore[method-assign]
|
||
{
|
||
"id": "m-1",
|
||
"address": "oc123@a.test",
|
||
"raw": "Subject: Your ChatGPT code is 123456",
|
||
"metadata": {"provider": "openai"},
|
||
}
|
||
]
|
||
|
||
codes = provider.poll_verification_codes(
|
||
apm.Mailbox(
|
||
email="oc123@a.test",
|
||
token="jwt",
|
||
domain="a.test",
|
||
failure_target="a.test",
|
||
),
|
||
seen_ids=set(),
|
||
)
|
||
self.assertEqual(codes, ["123456"])
|
||
|
||
def test_build_mail_provider_supports_cfmail(self):
|
||
provider = apm.build_mail_provider(
|
||
{
|
||
"mail": {"provider": "cfmail"},
|
||
"cfmail": {
|
||
"api_base": "https://mail.example.com",
|
||
"api_key": "pw",
|
||
"domains": ["a.test", "b.test"],
|
||
},
|
||
},
|
||
proxy="",
|
||
logger=logging.getLogger("test-build-cfmail"),
|
||
)
|
||
self.assertIsInstance(provider, apm.CfmailProvider)
|
||
self.assertEqual(provider.domains, ["a.test", "b.test"])
|
||
|
||
def test_api_server_merge_cfmail_api_key_preserves_masked_entries(self):
|
||
import api_server as aps
|
||
|
||
current = {
|
||
"cfmail": {
|
||
"api_base": "https://mail.example.com",
|
||
"api_key": "secret-1",
|
||
"domains": ["a.test"],
|
||
}
|
||
}
|
||
incoming = {
|
||
"cfmail": {
|
||
"api_base": "https://mail.example.com",
|
||
"api_key": aps.MASKED_VALUE,
|
||
"domains": ["a.test"],
|
||
}
|
||
}
|
||
|
||
merged = aps.merge_config_with_sensitive_fields(current, incoming)
|
||
self.assertEqual(merged["cfmail"]["api_key"], "secret-1")
|
||
|
||
class ProtocolRegistrarTests(unittest.TestCase):
|
||
def test_protocol_registrar_defaults_to_chatgpt_web_entry_mode(self):
|
||
logger = logging.getLogger("test-registration-default-entry-mode")
|
||
registrar = apm.ProtocolRegistrar(proxy="", logger=logger, conf={})
|
||
|
||
self.assertEqual(registrar.entry_mode, "chatgpt_web")
|
||
self.assertEqual(registrar._entry_mode_candidates(), ["chatgpt_web", "direct_auth"])
|
||
|
||
def test_capture_registration_tokens_uses_consent_url_redirect_code(self):
|
||
logger = logging.getLogger("test-registration-consent-code")
|
||
registrar = apm.ProtocolRegistrar(proxy="", logger=logger, conf={"flow": {"step_retry_attempts": 1}})
|
||
access_token = build_test_jwt(
|
||
{
|
||
"email": "jwt@example.com",
|
||
"exp": 1760000000,
|
||
"https://api.openai.com/auth": {"chatgpt_account_id": "acct_123"},
|
||
}
|
||
)
|
||
|
||
class FakeSession:
|
||
def __init__(self):
|
||
self.cookies = []
|
||
self.calls = []
|
||
|
||
def get(self, url, **kwargs):
|
||
self.calls.append((url, kwargs))
|
||
if url == "https://auth.openai.com/sign-in-with-chatgpt/codex/consent":
|
||
response = DummyResponse(302)
|
||
response.headers = {
|
||
"Location": (
|
||
"http://localhost:1455/auth/callback"
|
||
"?code=oauth-consent-code"
|
||
"&scope=openid+email+profile+offline_access"
|
||
"&state=oauth-state"
|
||
)
|
||
}
|
||
response.url = url
|
||
return response
|
||
if (
|
||
url
|
||
== "https://chatgpt.com/api/auth/callback/openai"
|
||
"?code=oauth-consent-code&scope=openid+email+profile+offline_access&state=oauth-state"
|
||
):
|
||
response = DummyResponse(200)
|
||
response.url = url
|
||
return response
|
||
if url == "https://chatgpt.com/api/auth/session":
|
||
response = DummyResponse(200, payload={"accessToken": access_token, "user": {"email": "jwt@example.com"}})
|
||
response.url = url
|
||
return response
|
||
raise AssertionError(f"unexpected url: {url}")
|
||
|
||
registrar.session = FakeSession() # type: ignore[assignment]
|
||
|
||
registrar._capture_registration_tokens( # type: ignore[attr-defined]
|
||
{"continue_url": "https://auth.openai.com/sign-in-with-chatgpt/codex/consent"}
|
||
)
|
||
|
||
self.assertEqual(registrar.registration_auth_code, "oauth-consent-code")
|
||
self.assertIsNotNone(registrar.registration_tokens)
|
||
self.assertEqual(registrar.registration_tokens["access_token"], access_token)
|
||
self.assertEqual(registrar.registration_tokens["email"], "jwt@example.com")
|
||
|
||
def test_capture_registration_tokens_falls_back_to_default_consent_when_add_phone_has_no_code(self):
|
||
logger = logging.getLogger("test-registration-add-phone-fallback")
|
||
registrar = apm.ProtocolRegistrar(proxy="", logger=logger, conf={"flow": {"step_retry_attempts": 1}})
|
||
access_token = build_test_jwt(
|
||
{
|
||
"email": "jwt@example.com",
|
||
"exp": 1760000000,
|
||
"https://api.openai.com/auth": {"chatgpt_account_id": "acct_123"},
|
||
}
|
||
)
|
||
|
||
class FakeSession:
|
||
def __init__(self):
|
||
self.cookies = []
|
||
self.calls = []
|
||
self.callback_completed = False
|
||
|
||
def get(self, url, **kwargs):
|
||
self.calls.append((url, kwargs))
|
||
if url == "https://auth.openai.com/add-phone":
|
||
response = DummyResponse(200, payload={"continue_url": "https://auth.openai.com/add-phone"})
|
||
response.url = url
|
||
return response
|
||
if url == "https://auth.openai.com/sign-in-with-chatgpt/codex/consent":
|
||
response = DummyResponse(302)
|
||
response.headers = {
|
||
"Location": (
|
||
"http://localhost:1455/auth/callback"
|
||
"?code=oauth-consent-code"
|
||
"&scope=openid+email+profile+offline_access"
|
||
"&state=oauth-state"
|
||
)
|
||
}
|
||
response.url = url
|
||
return response
|
||
if (
|
||
url
|
||
== "https://chatgpt.com/api/auth/callback/openai"
|
||
"?code=oauth-consent-code&scope=openid+email+profile+offline_access&state=oauth-state"
|
||
):
|
||
self.callback_completed = True
|
||
response = DummyResponse(200)
|
||
response.url = url
|
||
return response
|
||
if url == "https://chatgpt.com/api/auth/session":
|
||
payload = {"accessToken": access_token, "user": {"email": "jwt@example.com"}} if self.callback_completed else {}
|
||
response = DummyResponse(200, payload=payload)
|
||
response.url = url
|
||
return response
|
||
raise AssertionError(f"unexpected url: {url}")
|
||
|
||
registrar.session = FakeSession() # type: ignore[assignment]
|
||
|
||
registrar._capture_registration_tokens( # type: ignore[attr-defined]
|
||
{"continue_url": "https://auth.openai.com/add-phone"}
|
||
)
|
||
|
||
self.assertEqual(registrar.registration_auth_code, "oauth-consent-code")
|
||
self.assertIsNotNone(registrar.registration_tokens)
|
||
self.assertEqual(registrar.registration_tokens["access_token"], access_token)
|
||
self.assertIn(
|
||
"https://auth.openai.com/sign-in-with-chatgpt/codex/consent",
|
||
[call[0] for call in registrar.session.calls],
|
||
)
|
||
|
||
def test_capture_registration_tokens_uses_nested_create_account_code_without_following_consent(self):
|
||
logger = logging.getLogger("test-registration-nested-create-account-code")
|
||
registrar = apm.ProtocolRegistrar(proxy="", logger=logger, conf={"flow": {"step_retry_attempts": 1}})
|
||
access_token = build_test_jwt(
|
||
{
|
||
"email": "jwt@example.com",
|
||
"exp": 1760000000,
|
||
"https://api.openai.com/auth": {"chatgpt_account_id": "acct_123"},
|
||
}
|
||
)
|
||
|
||
class FakeSession:
|
||
def __init__(self):
|
||
self.cookies = []
|
||
self.calls = []
|
||
|
||
def get(self, url, **kwargs):
|
||
self.calls.append((url, kwargs))
|
||
if (
|
||
url
|
||
== "https://chatgpt.com/api/auth/callback/openai"
|
||
"?code=oauth-create-account-code&scope=openid+email+profile+offline_access&state=oauth-state"
|
||
):
|
||
response = DummyResponse(200)
|
||
response.url = url
|
||
return response
|
||
if url == "https://chatgpt.com/api/auth/session":
|
||
response = DummyResponse(200, payload={"accessToken": access_token, "user": {"email": "jwt@example.com"}})
|
||
response.url = url
|
||
return response
|
||
raise AssertionError(f"unexpected url: {url}")
|
||
|
||
registrar.session = FakeSession() # type: ignore[assignment]
|
||
|
||
registrar._capture_registration_tokens( # type: ignore[attr-defined]
|
||
{
|
||
"continue_url": "https://auth.openai.com/add-phone",
|
||
"page": {"type": "add_phone"},
|
||
"data": {
|
||
"oauth_callback": {
|
||
"code": "oauth-create-account-code",
|
||
"scope": "openid email profile offline_access",
|
||
"state": "oauth-state",
|
||
}
|
||
},
|
||
}
|
||
)
|
||
|
||
self.assertEqual(registrar.registration_auth_code, "oauth-create-account-code")
|
||
self.assertIsNotNone(registrar.registration_tokens)
|
||
self.assertEqual(registrar.registration_tokens["access_token"], access_token)
|
||
self.assertEqual(
|
||
[call[0] for call in registrar.session.calls],
|
||
[
|
||
(
|
||
"https://chatgpt.com/api/auth/callback/openai"
|
||
"?code=oauth-create-account-code&scope=openid+email+profile+offline_access&state=oauth-state"
|
||
),
|
||
"https://chatgpt.com/api/auth/session",
|
||
],
|
||
)
|
||
|
||
def test_capture_registration_tokens_uses_session_cookie_callback_without_following_consent(self):
|
||
logger = logging.getLogger("test-registration-cookie-callback-code")
|
||
registrar = apm.ProtocolRegistrar(proxy="", logger=logger, conf={"flow": {"step_retry_attempts": 1}})
|
||
access_token = build_test_jwt(
|
||
{
|
||
"email": "jwt@example.com",
|
||
"exp": 1760000000,
|
||
"https://api.openai.com/auth": {"chatgpt_account_id": "acct_123"},
|
||
}
|
||
)
|
||
cookie_payload = base64.urlsafe_b64encode(
|
||
json.dumps(
|
||
{
|
||
"continue_url": (
|
||
"http://localhost:1455/auth/callback"
|
||
"?code=oauth-cookie-code"
|
||
"&scope=openid+email+profile+offline_access"
|
||
"&state=oauth-state"
|
||
)
|
||
}
|
||
).encode("utf-8")
|
||
).rstrip(b"=").decode("ascii")
|
||
|
||
class DummyCookie:
|
||
def __init__(self, name, value):
|
||
self.name = name
|
||
self.value = value
|
||
self.domain = ".auth.openai.com"
|
||
self.path = "/"
|
||
|
||
class FakeSession:
|
||
def __init__(self):
|
||
self.cookies = [DummyCookie("oai-client-auth-session-info", f"{cookie_payload}.sig")]
|
||
self.calls = []
|
||
|
||
def get(self, url, **kwargs):
|
||
self.calls.append((url, kwargs))
|
||
if (
|
||
url
|
||
== "https://chatgpt.com/api/auth/callback/openai"
|
||
"?code=oauth-cookie-code&scope=openid+email+profile+offline_access&state=oauth-state"
|
||
):
|
||
response = DummyResponse(200)
|
||
response.url = url
|
||
return response
|
||
if url == "https://chatgpt.com/api/auth/session":
|
||
response = DummyResponse(200, payload={"accessToken": access_token, "user": {"email": "jwt@example.com"}})
|
||
response.url = url
|
||
return response
|
||
raise AssertionError(f"unexpected url: {url}")
|
||
|
||
registrar.session = FakeSession() # type: ignore[assignment]
|
||
|
||
registrar._capture_registration_tokens( # type: ignore[attr-defined]
|
||
{
|
||
"continue_url": "https://auth.openai.com/add-phone",
|
||
"page": {"type": "add_phone"},
|
||
}
|
||
)
|
||
|
||
self.assertEqual(registrar.registration_auth_code, "oauth-cookie-code")
|
||
self.assertIsNotNone(registrar.registration_tokens)
|
||
self.assertEqual(registrar.registration_tokens["access_token"], access_token)
|
||
self.assertEqual(
|
||
[call[0] for call in registrar.session.calls],
|
||
[
|
||
(
|
||
"https://chatgpt.com/api/auth/callback/openai"
|
||
"?code=oauth-cookie-code&scope=openid+email+profile+offline_access&state=oauth-state"
|
||
),
|
||
"https://chatgpt.com/api/auth/session",
|
||
],
|
||
)
|
||
|
||
def test_step4_validate_otp_sentinel_fallback(self):
|
||
logger = logging.getLogger("test-step4")
|
||
conf = {
|
||
"flow": {
|
||
"step_retry_attempts": 1,
|
||
"register_otp_validate_order": "normal,sentinel",
|
||
}
|
||
}
|
||
registrar = apm.ProtocolRegistrar(proxy="", logger=logger, conf=conf)
|
||
registrar.sentinel_gen.generate_token = lambda *_args, **_kwargs: "token-sentinel"
|
||
|
||
captured_headers = []
|
||
|
||
def fake_post(_url, **kwargs):
|
||
captured_headers.append(kwargs.get("headers") or {})
|
||
if len(captured_headers) == 1:
|
||
return DummyResponse(400)
|
||
return DummyResponse(200)
|
||
|
||
registrar.session.post = fake_post
|
||
|
||
ok = registrar.step4_validate_otp("123456")
|
||
|
||
self.assertTrue(ok)
|
||
self.assertEqual(len(captured_headers), 2)
|
||
self.assertNotIn("openai-sentinel-token", captured_headers[0])
|
||
self.assertEqual(captured_headers[1].get("openai-sentinel-token"), "token-sentinel")
|
||
|
||
def test_register_passes_mail_poll_interval_to_provider(self):
|
||
logger = logging.getLogger("test-register-mail-poll-interval")
|
||
registrar = apm.ProtocolRegistrar(proxy="", logger=logger, conf={"flow": {"step_retry_attempts": 1}})
|
||
|
||
registrar.step0_init_oauth_session = lambda *_args, **_kwargs: True
|
||
registrar.step2_register_user = lambda *_args, **_kwargs: True
|
||
registrar.step3_send_otp = lambda *_args, **_kwargs: True
|
||
registrar.step4_validate_otp = lambda *_args, **_kwargs: True
|
||
registrar.step5_create_account = lambda *_args, **_kwargs: True
|
||
|
||
class FakeMailProvider:
|
||
provider_name = "fake"
|
||
|
||
def __init__(self):
|
||
self.called_kwargs = {}
|
||
|
||
def wait_for_verification_code(self, _mailbox, **kwargs):
|
||
self.called_kwargs = kwargs
|
||
return "123456"
|
||
|
||
provider = FakeMailProvider()
|
||
|
||
with patch("auto_pool_maintainer.time.sleep", lambda *_args, **_kwargs: None):
|
||
ok = registrar.register(
|
||
email="test@example.com",
|
||
password="pw",
|
||
client_id="cid",
|
||
redirect_uri="http://localhost/cb",
|
||
mailbox=apm.Mailbox(email="test@example.com"),
|
||
mail_provider=provider, # type: ignore[arg-type]
|
||
otp_timeout_seconds=88,
|
||
otp_poll_interval_seconds=1.25,
|
||
)
|
||
|
||
self.assertTrue(ok)
|
||
self.assertEqual(provider.called_kwargs.get("timeout"), 88)
|
||
self.assertEqual(provider.called_kwargs.get("poll_interval_seconds"), 1.25)
|
||
|
||
|
||
class RegisterOneFlowTests(unittest.TestCase):
|
||
class _FakeMailProvider:
|
||
provider_name = "fake"
|
||
|
||
@staticmethod
|
||
def create_mailbox():
|
||
return apm.Mailbox(email="fake@example.com")
|
||
|
||
class _FakeRuntime:
|
||
def __init__(self, oauth_token=None):
|
||
self.stop_event = threading.Event()
|
||
self.target_tokens = 1
|
||
self._token_count = 0
|
||
self.mail_provider = RegisterOneFlowTests._FakeMailProvider()
|
||
self.mail_provider_name = "fake"
|
||
self.logger = logging.getLogger("test-register-one")
|
||
self.proxy = ""
|
||
self.conf = {}
|
||
self.oauth_client_id = "cid"
|
||
self.oauth_redirect_uri = "http://localhost/cb"
|
||
self.mail_otp_timeout_seconds = 60
|
||
self.mail_poll_interval_seconds = 1.0
|
||
self.oauth_outer_retry_attempts = 3
|
||
self.last_oauth_failure_detail = ""
|
||
self.oauth_token = oauth_token
|
||
self.oauth_called = False
|
||
self.saved_tokens = None
|
||
self.saved_account = None
|
||
self.success_key = None
|
||
|
||
def get_token_success_count(self):
|
||
return self._token_count
|
||
|
||
def wait_for_provider_availability(self, worker_id=0):
|
||
return None
|
||
|
||
def oauth_login_with_retry(self, mailbox, password):
|
||
self.oauth_called = True
|
||
return self.oauth_token
|
||
|
||
def claim_token_slot(self):
|
||
self._token_count += 1
|
||
return True, self._token_count
|
||
|
||
def release_token_slot(self):
|
||
self._token_count = max(0, self._token_count - 1)
|
||
|
||
def save_tokens(self, email, tokens):
|
||
self.saved_tokens = tokens
|
||
return True
|
||
|
||
def save_account(self, email, password):
|
||
self.saved_account = (email, password)
|
||
|
||
def note_attempt_success(self, success_key="register_oauth_success"):
|
||
self.success_key = success_key
|
||
|
||
def note_attempt_failure(self, stage, email="", detail=""):
|
||
raise AssertionError(f"unexpected failure: stage={stage} email={email} detail={detail}")
|
||
|
||
class _FakeRegistrar:
|
||
def __init__(self, proxy, logger, conf):
|
||
self.last_failure_detail = ""
|
||
self.last_failure_stage = ""
|
||
|
||
def register(self, **kwargs):
|
||
return True
|
||
|
||
def exchange_codex_tokens(self, client_id, redirect_uri):
|
||
raise AssertionError("register_one 不应再调用 exchange_codex_tokens")
|
||
|
||
def test_register_one_calls_oauth_path(self):
|
||
fake_runtime = self._FakeRuntime(oauth_token={"access_token": "oauth-token"})
|
||
|
||
class Registrar(self._FakeRegistrar):
|
||
pass
|
||
|
||
with patch("auto_pool_maintainer.ProtocolRegistrar", Registrar), patch(
|
||
"auto_pool_maintainer.generate_random_password", lambda: "Pw123456!"
|
||
):
|
||
_, success, _, _ = apm.register_one(fake_runtime, worker_id=1)
|
||
|
||
self.assertTrue(success)
|
||
self.assertTrue(fake_runtime.oauth_called)
|
||
self.assertEqual(fake_runtime.saved_tokens, {"access_token": "oauth-token"})
|
||
self.assertEqual(fake_runtime.success_key, "register_oauth_success")
|
||
|
||
def test_register_one_prefers_registration_session_tokens(self):
|
||
class RuntimeWithoutOauth(self._FakeRuntime):
|
||
def oauth_login_with_retry(self, mailbox, password):
|
||
raise AssertionError("已有注册阶段 token 时不应再跑 OAuth 登录")
|
||
|
||
runtime = RuntimeWithoutOauth(oauth_token=None)
|
||
|
||
class Registrar(self._FakeRegistrar):
|
||
def __init__(self, proxy, logger, conf):
|
||
super().__init__(proxy, logger, conf)
|
||
self.registration_tokens = {"access_token": "session-token", "email": "fake@example.com"}
|
||
|
||
with patch("auto_pool_maintainer.ProtocolRegistrar", Registrar), patch(
|
||
"auto_pool_maintainer.generate_random_password", lambda: "Pw123456!"
|
||
):
|
||
_, success, _, _ = apm.register_one(runtime, worker_id=1)
|
||
|
||
self.assertTrue(success)
|
||
self.assertEqual(runtime.saved_tokens, {"access_token": "session-token", "email": "fake@example.com"})
|
||
self.assertEqual(runtime.success_key, "register_oauth_success")
|
||
|
||
def test_register_one_returns_fail_when_oauth_failed(self):
|
||
class RuntimeWithFailure(self._FakeRuntime):
|
||
failure_events = []
|
||
|
||
def note_attempt_failure(self, stage, email="", detail=""):
|
||
self.failure_events.append((stage, email, detail))
|
||
|
||
runtime = RuntimeWithFailure(oauth_token=None)
|
||
|
||
class Registrar(self._FakeRegistrar):
|
||
pass
|
||
|
||
with patch("auto_pool_maintainer.ProtocolRegistrar", Registrar), patch(
|
||
"auto_pool_maintainer.generate_random_password", lambda: "Pw123456!"
|
||
):
|
||
_, success, _, _ = apm.register_one(runtime, worker_id=1)
|
||
|
||
self.assertFalse(success)
|
||
self.assertTrue(runtime.oauth_called)
|
||
self.assertTrue(runtime.failure_events)
|
||
self.assertEqual(runtime.failure_events[-1][0], "oauth")
|
||
|
||
def test_register_one_create_mailbox_failure_marks_selected_domain(self):
|
||
class FakeMailProvider:
|
||
provider_name = "fake"
|
||
|
||
def __init__(self):
|
||
self.last_selected_domain = "a.test"
|
||
self.failure_calls = []
|
||
|
||
def wait_for_availability(self, worker_id=0):
|
||
return None
|
||
|
||
def create_mailbox(self):
|
||
return None
|
||
|
||
def note_domain_failure(self, domain, *, stage, detail=""):
|
||
self.failure_calls.append((domain, stage, detail))
|
||
|
||
def note_domain_success(self, domain):
|
||
return None
|
||
|
||
class FakeRuntime(self._FakeRuntime):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.mail_provider = FakeMailProvider()
|
||
self.mail_provider_name = "fake"
|
||
self.failure_events = []
|
||
|
||
def note_attempt_failure(self, stage, email="", detail=""):
|
||
self.failure_events.append((stage, email, detail))
|
||
|
||
runtime = FakeRuntime()
|
||
email, success, _, _ = apm.register_one(runtime)
|
||
|
||
self.assertIsNone(email)
|
||
self.assertFalse(success)
|
||
self.assertEqual(runtime.mail_provider.failure_calls, [("a.test", "create_mailbox", "provider=fake")])
|
||
|
||
def test_register_one_register_mail_timeout_marks_mailbox_domain(self):
|
||
class FakeMailProvider(self._FakeMailProvider):
|
||
provider_name = "fake"
|
||
|
||
def __init__(self):
|
||
self.failure_calls = []
|
||
|
||
def wait_for_availability(self, worker_id=0):
|
||
return None
|
||
|
||
@staticmethod
|
||
def create_mailbox():
|
||
return apm.Mailbox(email="fake@example.com", domain="a.test")
|
||
|
||
def note_domain_failure(self, domain, *, stage, detail=""):
|
||
self.failure_calls.append((domain, stage, detail))
|
||
|
||
def note_domain_success(self, domain):
|
||
return None
|
||
|
||
class FakeRuntime(self._FakeRuntime):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.mail_provider = FakeMailProvider()
|
||
self.mail_provider_name = "fake"
|
||
self.failure_events = []
|
||
|
||
def note_attempt_failure(self, stage, email="", detail=""):
|
||
self.failure_events.append((stage, email, detail))
|
||
|
||
class FakeRegistrar(self._FakeRegistrar):
|
||
def register(self, **kwargs):
|
||
self.last_failure_stage = "register_mail_otp_timeout"
|
||
self.last_failure_detail = "provider=fake"
|
||
return False
|
||
|
||
runtime = FakeRuntime()
|
||
with patch("auto_pool_maintainer.ProtocolRegistrar", FakeRegistrar), patch(
|
||
"auto_pool_maintainer.generate_random_password", lambda: "Pw123456!"
|
||
):
|
||
email, success, _, _ = apm.register_one(runtime, worker_id=1)
|
||
|
||
self.assertEqual(email, "fake@example.com")
|
||
self.assertFalse(success)
|
||
self.assertEqual(runtime.mail_provider.failure_calls, [("a.test", "register", "provider=fake")])
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|