From a833cf429e042efe09cf0991a5b8b1daca24d4df Mon Sep 17 00:00:00 2001
From: 9qeklajc <tthefux@gmail.com>
Date: Sun, 29 Mar 2026 21:22:09 +0200
Subject: [PATCH] add refund x-cashu to default refund endpoint

---
 routstr/balance.py         |  19 +++++++
 tests/unit/test_balance.py | 100 +++++++++++++++++++++++++++++++++++++
 2 files changed, 119 insertions(+)
 create mode 100644 tests/unit/test_balance.py

diff --git a/routstr/balance.py b/routstr/balance.py
index 9adbd9b7..6eed34aa 100644
--- a/routstr/balance.py
+++ b/routstr/balance.py
@@ -5,6 +5,7 @@ from time import monotonic
 from typing import Annotated, NoReturn
 
 from fastapi import APIRouter, Depends, Header, HTTPException
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 from sqlmodel import select
 
@@ -207,6 +208,7 @@ async def _refund_cache_set(authorization: str, value: dict[str, str]) -> None:
 @router.post("/refund")
 async def refund_wallet_endpoint(
     authorization: Annotated[str, Header(...)],
+    x_cashu: Annotated[str | None, Header()] = None,
     session: AsyncSession = Depends(get_session),
 ) -> dict[str, str]:
     if not authorization.startswith("Bearer "):
@@ -217,6 +219,23 @@ async def refund_wallet_endpoint(
 
     bearer_value: str = authorization[7:]
 
+    if x_cashu:
+        payment_token_hash = hashlib.sha256(x_cashu.strip().encode()).hexdigest()
+        result = await session.get(CashuTransaction, payment_token_hash)
+        if result is None:
+            raise HTTPException(status_code=404, detail="Refund not found")
+        if result.swept:
+            raise HTTPException(status_code=410, detail="Refund has been swept")
+        result.collected = True
+        session.add(result)
+        await session.commit()
+        body: dict[str, str] = {"token": result.token}
+        if result.unit == "sat":
+            body["sats"] = str(result.amount)
+        else:
+            body["msats"] = str(result.amount)
+        return JSONResponse(content=body, headers={"X-Cashu": result.token})
+
     key: ApiKey = await validate_bearer_key(bearer_value, session)
 
     if key.total_balance <= 0:
diff --git a/tests/unit/test_balance.py b/tests/unit/test_balance.py
new file mode 100644
index 00000000..fdb4d0fc
--- /dev/null
+++ b/tests/unit/test_balance.py
@@ -0,0 +1,100 @@
+import hashlib
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from routstr.balance import refund_wallet_endpoint
+from routstr.core.db import CashuTransaction
+
+
+def _make_cashu_tx(token: str, amount: int, unit: str, swept: bool = False) -> CashuTransaction:
+    tx = CashuTransaction(token=token, amount=amount, unit=unit)
+    tx.swept = swept
+    tx.collected = False
+    return tx
+
+
+@pytest.mark.asyncio
+async def test_refund_x_cashu_returns_token() -> None:
+    x_cashu_token = "cashuAtest_token_value"
+    expected_hash = hashlib.sha256(x_cashu_token.strip().encode()).hexdigest()
+    tx = _make_cashu_tx(token="cashuArefund_token", amount=1000, unit="msat")
+
+    session = MagicMock()
+    session.get = AsyncMock(return_value=tx)
+    session.add = MagicMock()
+    session.commit = AsyncMock()
+
+    result = await refund_wallet_endpoint(
+        authorization="Bearer sk-somekey",
+        x_cashu=x_cashu_token,
+        session=session,
+    )
+
+    session.get.assert_awaited_once_with(CashuTransaction, expected_hash)
+    import json
+    body = json.loads(result.body)
+    assert body["token"] == "cashuArefund_token"
+    assert body["msats"] == "1000"
+    assert result.headers["X-Cashu"] == "cashuArefund_token"
+    assert tx.collected is True
+
+
+@pytest.mark.asyncio
+async def test_refund_x_cashu_sat_unit() -> None:
+    x_cashu_token = "cashuAsat_token"
+    tx = _make_cashu_tx(token="cashuArefund_sat", amount=500, unit="sat")
+
+    session = MagicMock()
+    session.get = AsyncMock(return_value=tx)
+    session.add = MagicMock()
+    session.commit = AsyncMock()
+
+    result = await refund_wallet_endpoint(
+        authorization="Bearer sk-somekey",
+        x_cashu=x_cashu_token,
+        session=session,
+    )
+
+    import json
+    body = json.loads(result.body)
+    assert body["token"] == "cashuArefund_sat"
+    assert body["sats"] == "500"
+    assert "msats" not in body
+    assert result.headers["X-Cashu"] == "cashuArefund_sat"
+
+
+@pytest.mark.asyncio
+async def test_refund_x_cashu_not_found_raises_404() -> None:
+    from fastapi import HTTPException
+
+    session = MagicMock()
+    session.get = AsyncMock(return_value=None)
+
+    with pytest.raises(HTTPException) as exc_info:
+        await refund_wallet_endpoint(
+            authorization="Bearer sk-somekey",
+            x_cashu="cashuAmissing_token",
+            session=session,
+        )
+
+    assert exc_info.value.status_code == 404
+
+
+@pytest.mark.asyncio
+async def test_refund_x_cashu_swept_raises_410() -> None:
+    from fastapi import HTTPException
+
+    tx = _make_cashu_tx(token="cashuAswept", amount=100, unit="msat", swept=True)
+
+    session = MagicMock()
+    session.get = AsyncMock(return_value=tx)
+
+    with pytest.raises(HTTPException) as exc_info:
+        await refund_wallet_endpoint(
+            authorization="Bearer sk-somekey",
+            x_cashu="cashuAswept_token",
+            session=session,
+        )
+
+    assert exc_info.value.status_code == 410