Enhance Selectel balance handling and error messaging in GPU rental system

- Updated Selectel documentation to clarify error responses related to balance issues. - Introduced `peek_balance_rub` function for best-effort balance checks, handling API failures gracefully. - Improved error messages for insufficient funds and quota issues, specifying Selectel's 403 policy response. - Added balance checks in the doctor command to ensure users are informed about their balance status before attempting GPU creation. - Refactored exception handling in cloud operations to provide clearer feedback on balance-related errors. Co-authored-by: Cursor <cursoragent@cursor.com>
2026-08-23 17:54:11 +03:00
parent 322577cf9f
commit 0bad693b12
10 changed files with 389 additions and 28 deletions
@@ -1,9 +1,11 @@
 from gpu_rent.balance import (
    BalanceWatchState,
    balance_rub_from_payload,
+    peek_balance_rub,
    plan_balance_notices,
    spend_steps,
 )
+from gpu_rent.errors import CloudError


 def test_balance_kopecks():
@@ -54,3 +56,17 @@ def test_plan_low_once():
    assert any("низкий" in m for m in msgs)
    state, msgs2 = plan_balance_notices(state, 100.0, step_rub=200.0, low_rub=200.0)
    assert not any("низкий" in m for m in msgs2)
+
+
+def test_peek_balance_rub_empty_token():
+    assert peek_balance_rub("") is None
+    assert peek_balance_rub(None) is None
+
+
+def test_peek_balance_rub_swallows_errors(monkeypatch):
+    def boom(*_a, **_k):
+        raise CloudError("сеть")
+
+    monkeypatch.setattr("gpu_rent.balance.fetch_balance_rub", boom)
+    assert peek_balance_rub("tok") is None
+
@@ -0,0 +1,111 @@
+from types import SimpleNamespace
+
+import pytest
+
+from gpu_rent.cloud import (
+    _exc_brief,
+    _http_status,
+    _is_create_policy_deny,
+    _wrap,
+    create_gpu_server,
+)
+from gpu_rent.errors import CloudError
+
+POLICY_403 = (
+    "ForbiddenException: 403: Client Error for url: "
+    "https://ru-6.cloud.api.selcloud.ru/compute/v2.1/servers, "
+    "Policy doesn't allow os_compute_api:servers:create to be performed."
+)
+
+
+class Forbidden(Exception):
+    def __init__(self, msg: str, http_status: int = 403):
+        super().__init__(msg)
+        self.http_status = http_status
+
+
+def test_http_status_from_attr_and_text():
+    assert _http_status(Forbidden(POLICY_403)) == 403
+    assert _http_status(RuntimeError("no status")) is None
+
+
+def test_is_create_policy_deny():
+    assert _is_create_policy_deny(Forbidden(POLICY_403))
+    assert not _is_create_policy_deny(RuntimeError("quota exceeded for cores"))
+
+
+def test_exc_brief_strips_url_and_user_data():
+    dump = (
+        "BadRequestException: 400: Invalid input for field/attribute user_data. "
+        "Value: #cloud-config\nssh_authorized_keys:\n"
+        "  - ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIFake gpu-rent\n"
+    )
+    brief = _exc_brief(RuntimeError(dump))
+    assert "ssh-ed25519" not in brief
+    assert "AAAAC3" not in brief
+    assert "http" not in brief.lower() or "https://" not in brief
+
+
+def test_wrap_403_create_is_funds_not_gpu_quota():
+    err = _wrap(Forbidden(POLICY_403), "create server")
+    msg = str(err)
+    assert isinstance(err, CloudError)
+    assert "GPU=0" not in msg
+    assert "тикет" not in msg
+    assert "ssh-ed25519" not in msg
+    assert "https://" not in msg
+    assert "Пополни" in msg or "нет средств" in msg or "кончились деньги" in msg
+
+
+def test_wrap_403_create_with_zero_balance(monkeypatch):
+    monkeypatch.setattr("gpu_rent.balance.peek_balance_rub", lambda *_a, **_k: 0.0)
+    msg = str(_wrap(Forbidden(POLICY_403), "create server", api_token="tok"))
+    assert "0 ₽" in msg
+    assert "нет средств" in msg
+    assert "Пополни" in msg
+    assert "GPU=0" not in msg
+
+
+def test_wrap_403_create_with_healthy_balance_is_iam(monkeypatch):
+    monkeypatch.setattr("gpu_rent.balance.peek_balance_rub", lambda *_a, **_k: 1500.0)
+    msg = str(_wrap(Forbidden(POLICY_403), "create server", api_token="tok"))
+    assert "1500 ₽" in msg
+    assert "IAM" in msg
+    assert "GPU=0" not in msg
+
+
+def test_wrap_quota_message_only_when_quota_in_error():
+    exc = RuntimeError("403 OverLimit: quota exceeded for instances")
+    msg = str(_wrap(exc, "create server"))
+    assert "квота" in msg.lower()
+
+
+def test_create_gpu_server_stops_on_first_403():
+    calls: list[int] = []
+
+    class Compute:
+        def create_server(self, **kw):
+            calls.append(1)
+            raise Forbidden(POLICY_403)
+
+    logs: list[str] = []
+    with pytest.raises(CloudError) as caught:
+        create_gpu_server(
+            SimpleNamespace(compute=Compute()),
+            flavor_id="1",
+            net_id="n",
+            sg_name="g",
+            boot_volume_id="b",
+            data_volume_id="d",
+            az="ru-6a",
+            spot=True,
+            public_key="ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIFake gpu-rent",
+            log=logs.append,
+        )
+    assert len(calls) == 1
+    assert logs == []
+    msg = str(caught.value)
+    assert "GPU=0" not in msg
+    assert "тикет" not in msg
+    assert "ssh-ed25519" not in msg
+    assert "Пополни" in msg or "нет средств" in msg or "кончились деньги" in msg
@@ -0,0 +1,51 @@
+from types import SimpleNamespace
+
+from gpu_rent.doctor import Check, _check_balance
+from gpu_rent.errors import CloudError
+
+
+def test_check_balance_skips_without_token():
+    checks: list[Check] = []
+    _check_balance(SimpleNamespace(selectel_api_token=""), checks)
+    assert checks[0].name == "баланс ₽"
+    assert checks[0].ok is True
+    assert checks[0].blocking is False
+
+
+def test_check_balance_blocks_when_empty(monkeypatch):
+    monkeypatch.setattr("gpu_rent.balance.fetch_balance_rub", lambda *_a, **_k: 0.0)
+    checks: list[Check] = []
+    _check_balance(SimpleNamespace(selectel_api_token="tok"), checks)
+    assert checks[0].ok is False
+    assert checks[0].blocking is True
+    assert "0 ₽" in checks[0].detail
+    assert "Пополни" in checks[0].detail
+
+
+def test_check_balance_warns_when_low(monkeypatch):
+    monkeypatch.setattr("gpu_rent.balance.fetch_balance_rub", lambda *_a, **_k: 40.0)
+    checks: list[Check] = []
+    _check_balance(SimpleNamespace(selectel_api_token="tok"), checks)
+    assert checks[0].ok is True
+    assert checks[0].blocking is False
+    assert "40 ₽" in checks[0].detail
+
+
+def test_check_balance_ok_when_healthy(monkeypatch):
+    monkeypatch.setattr("gpu_rent.balance.fetch_balance_rub", lambda *_a, **_k: 1200.0)
+    checks: list[Check] = []
+    _check_balance(SimpleNamespace(selectel_api_token="tok"), checks)
+    assert checks[0].ok is True
+    assert "1200 ₽" in checks[0].detail
+
+
+def test_check_balance_fetch_error_is_warning(monkeypatch):
+    def boom(*_a, **_k):
+        raise CloudError("сеть")
+
+    monkeypatch.setattr("gpu_rent.balance.fetch_balance_rub", boom)
+    checks: list[Check] = []
+    _check_balance(SimpleNamespace(selectel_api_token="tok"), checks)
+    assert checks[0].ok is True
+    assert checks[0].blocking is False
+    assert "не прочитать" in checks[0].detail
@@ -3,7 +3,14 @@ from types import SimpleNamespace
 from gpu_rent.cli import _stop_after_failed_up


-def test_stop_after_failed_up_calls_cmd_stop(monkeypatch):
+def _quiet(monkeypatch):
+    monkeypatch.setattr("gpu_rent.cli.warn", lambda *_a, **_k: None)
+    monkeypatch.setattr("gpu_rent.cli.ok", lambda *_a, **_k: None)
+    monkeypatch.setattr("gpu_rent.cli.err", lambda *_a, **_k: None)
+    monkeypatch.setattr("gpu_rent.cli.log", lambda *_a, **_k: None)
+
+
+def test_stop_after_failed_up_skips_when_no_server(monkeypatch):
    calls: list[dict] = []

    def fake_stop(cfg, *, no_pull=False, log=None, destroy_disks=False):
@@ -11,13 +18,28 @@ def test_stop_after_failed_up_calls_cmd_stop(monkeypatch):
        return SimpleNamespace()

    monkeypatch.setattr("gpu_rent.cli.cmd_stop", fake_stop)
-    monkeypatch.setattr("gpu_rent.cli.warn", lambda *_a, **_k: None)
-    monkeypatch.setattr("gpu_rent.cli.ok", lambda *_a, **_k: None)
-    monkeypatch.setattr("gpu_rent.cli.err", lambda *_a, **_k: None)
-    monkeypatch.setattr("gpu_rent.cli.log", lambda *_a, **_k: None)
+    _quiet(monkeypatch)
    monkeypatch.setattr(
        "gpu_rent.state.load_state",
-        lambda: SimpleNamespace(floating_ip=None),
+        lambda: SimpleNamespace(server_id=None, floating_ip=None),
+    )
+
+    _stop_after_failed_up(SimpleNamespace(), RuntimeError("boom"))
+    assert calls == []
+
+
+def test_stop_after_failed_up_runs_when_server_exists(monkeypatch):
+    calls: list[dict] = []
+
+    def fake_stop(cfg, *, no_pull=False, log=None, destroy_disks=False):
+        calls.append({"no_pull": no_pull, "destroy_disks": destroy_disks})
+        return SimpleNamespace()
+
+    monkeypatch.setattr("gpu_rent.cli.cmd_stop", fake_stop)
+    _quiet(monkeypatch)
+    monkeypatch.setattr(
+        "gpu_rent.state.load_state",
+        lambda: SimpleNamespace(server_id="s1", floating_ip=None),
    )

    _stop_after_failed_up(SimpleNamespace(), RuntimeError("boom"))
@@ -33,20 +55,20 @@ def test_stop_after_failed_up_prints_digest_when_fip(monkeypatch):
        return SimpleNamespace()

    monkeypatch.setattr("gpu_rent.cli.cmd_stop", fake_stop)
-    monkeypatch.setattr("gpu_rent.cli.warn", lambda *_a, **_k: None)
-    monkeypatch.setattr("gpu_rent.cli.ok", lambda *_a, **_k: None)
-    monkeypatch.setattr("gpu_rent.cli.err", lambda *_a, **_k: None)
-    monkeypatch.setattr("gpu_rent.cli.log", lambda *_a, **_k: None)
+    _quiet(monkeypatch)
    monkeypatch.setattr(
        "gpu_rent.state.load_state",
-        lambda: SimpleNamespace(floating_ip="10.0.0.1"),
+        lambda: SimpleNamespace(server_id="s1", floating_ip="10.0.0.1"),
    )
    monkeypatch.setattr(
        "gpu_rent.vm_logs.print_log_digest",
        lambda cfg, host, **kw: digests.append(host),
    )

-    _stop_after_failed_up(SimpleNamespace(enable_swarmui=True, llm_runtime="none"), RuntimeError("boom"))
+    _stop_after_failed_up(
+        SimpleNamespace(enable_swarmui=True, llm_runtime="none"),
+        RuntimeError("boom"),
+    )
    assert digests == ["10.0.0.1"]
    assert calls == [{"no_pull": True}]