From 6dd9dcb5d42e0660b8f634cc1a0ef5867db3902d Mon Sep 17 00:00:00 2001 From: bohd4nx Date: Mon, 11 May 2026 13:02:00 +0300 Subject: [PATCH] feat: update wallet handling to support separate TON and USDT balances - Refactor `get_wallet()` to return `ton_balance` and `usdt_balance` in `WalletInfo`. - Update transaction processing to validate balances for both TON and USDT payment methods. - Introduce `parse_required_payment_amount` utility to extract payment amounts from responses. - Modify tests to cover new balance checks and payment method handling. - Upgrade GitHub Actions artifact upload action to v7. - Enhance documentation and examples to reflect changes in wallet balance handling. --- .github/workflows/publish.yml | 2 +- CHANGELOG.md | 3 + README.md | 7 ++ examples/client/wallet_info.py | 5 +- pyfragment/client.py | 5 +- pyfragment/methods/giveaway_premium.py | 10 ++- pyfragment/methods/giveaway_stars.py | 10 ++- pyfragment/methods/purchase_premium.py | 10 ++- pyfragment/methods/purchase_stars.py | 10 ++- pyfragment/types/constants.py | 8 +- pyfragment/types/exceptions.py | 9 +- pyfragment/types/results.py | 8 +- pyfragment/utils/__init__.py | 3 +- pyfragment/utils/decoder.py | 16 ++-- pyfragment/utils/html.py | 15 ++++ pyfragment/utils/wallet.py | 119 ++++++++++++++++++++++--- tests/001_test_decode.py | 98 ++++++++++++++++++-- tests/003_test_balance.py | 46 +++++++++- tests/004_test_stars.py | 10 ++- tests/005_test_premium.py | 10 ++- tests/007_test_wallet.py | 12 ++- tests/015_test_payment_amount.py | 21 +++++ 22 files changed, 381 insertions(+), 56 deletions(-) create mode 100644 tests/015_test_payment_amount.py diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 2e97f59..9b9892e 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -56,7 +56,7 @@ jobs: - run: uv build - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@v7 with: name: dist path: dist/* diff --git a/CHANGELOG.md b/CHANGELOG.md index 98881cf..002a208 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,8 @@ and this project uses [Calendar Versioning](https://calver.org/) (`YYYY.MINOR.MI - Added runtime validation for `payment_method` via `SUPPORTED_PAYMENT_METHODS` and `ConfigurationError.INVALID_PAYMENT_METHOD` - Updated method docstrings to explicitly document recipient/channel formats: - `@username` / `username` / `https://t.me/username` +- `get_wallet()` now returns balances as separate fields: `ton_balance` and `usdt_balance` +- Wallet/system test output now prints TON and USDT balances on separate lines ### Tests @@ -29,6 +31,7 @@ and this project uses [Calendar Versioning](https://calver.org/) (`YYYY.MINOR.MI - invalid payment method - payment method propagation to `init*Request` payloads - accepted query formats (`@`, plain username, `t.me` link) +- Extended wallet tests to verify separate TON/USDT balance values in `WalletInfo` ### Documentation diff --git a/README.md b/README.md index c676e60..a762b0c 100644 --- a/README.md +++ b/README.md @@ -88,6 +88,9 @@ async def main() -> None: "stel_ton_token": "...", }, ) as client: + wallet = await client.get_wallet() + print(f"Wallet: {wallet.address} | TON: {wallet.ton_balance} | USDT: {wallet.usdt_balance}") + recipient = "https://t.me/username" # also supports: @username, username stars = await client.purchase_stars(recipient, amount=500, payment_method="usdt_ton") @@ -105,6 +108,10 @@ Full runnable examples: - https://github.com/bohd4nx/pyfragment/tree/master/examples - `examples/` +Payload debug/decode helper (thanks): + +- https://ton-cell-abi-viewer.vercel.app/ + ---
diff --git a/examples/client/wallet_info.py b/examples/client/wallet_info.py index 5ac628f..bfc6620 100644 --- a/examples/client/wallet_info.py +++ b/examples/client/wallet_info.py @@ -1,5 +1,5 @@ """ -Example: fetch wallet address, state, and balance. +Example: fetch wallet address, state, and separate TON/USDT balances. Cookies can be passed as a dict or as a JSON string. wallet_version defaults to "V5R1" — change to "V4R2" for older wallets. @@ -35,7 +35,8 @@ async def main() -> None: wallet = await client.get_wallet() print(f"Address: {wallet.address}") print(f"State: {wallet.state}") - print(f"Balance: {wallet.balance} TON") + print(f"Balance: {wallet.ton_balance} TON") + print(f"Balance: {wallet.usdt_balance} USDT") if __name__ == "__main__": diff --git a/pyfragment/client.py b/pyfragment/client.py index 2243957..992d033 100644 --- a/pyfragment/client.py +++ b/pyfragment/client.py @@ -193,11 +193,12 @@ class FragmentClient: return await recharge_ads(self, account, amount) async def get_wallet(self) -> WalletInfo: - """Return the address, state and balance of the TON wallet. + """Return the address, state, and balances of the wallet. Returns: :class:`WalletInfo` with ``address`` (``"UQ..."``), ``state`` - (``"active"``, ``"uninit"``, ``"nonexist"``, or ``"frozen"``), and ``balance`` in TON. + (``"active"``, ``"uninit"``, ``"nonexist"``, or ``"frozen"``), + ``ton_balance`` in TON, and ``usdt_balance`` in USDT. """ return await get_wallet_info(self) diff --git a/pyfragment/methods/giveaway_premium.py b/pyfragment/methods/giveaway_premium.py index c9b8a0d..32fc860 100644 --- a/pyfragment/methods/giveaway_premium.py +++ b/pyfragment/methods/giveaway_premium.py @@ -13,7 +13,7 @@ from pyfragment.types import ( VerificationError, ) from pyfragment.types.constants import DEVICE, PREMIUM_GIVEAWAY_PAGE, SUPPORTED_PAYMENT_METHODS, PaymentMethod -from pyfragment.utils import get_account_info, process_transaction +from pyfragment.utils import get_account_info, parse_required_payment_amount, process_transaction if TYPE_CHECKING: from pyfragment.client import FragmentClient @@ -77,6 +77,7 @@ async def giveaway_premium( }, page_url=PREMIUM_GIVEAWAY_PAGE, ) + required_payment_amount = parse_required_payment_amount(result, payment_method) req_id = result.get("req_id") if not req_id: raise FragmentAPIError(FragmentAPIError.NO_REQUEST_ID.format(context="Premium giveaway")) @@ -95,7 +96,12 @@ async def giveaway_premium( if transaction.get("need_verify"): raise VerificationError(VerificationError.KYC_REQUIRED) - tx_hash = await process_transaction(client, transaction) + tx_hash = await process_transaction( + client, + transaction, + payment_method=payment_method, + required_payment_amount=required_payment_amount, + ) return PremiumGiveawayResult( transaction_id=tx_hash, channel=channel, diff --git a/pyfragment/methods/giveaway_stars.py b/pyfragment/methods/giveaway_stars.py index 83369a8..8eac8b7 100644 --- a/pyfragment/methods/giveaway_stars.py +++ b/pyfragment/methods/giveaway_stars.py @@ -13,7 +13,7 @@ from pyfragment.types import ( VerificationError, ) from pyfragment.types.constants import DEVICE, STARS_GIVEAWAY_PAGE, SUPPORTED_PAYMENT_METHODS, PaymentMethod -from pyfragment.utils import get_account_info, process_transaction +from pyfragment.utils import get_account_info, parse_required_payment_amount, process_transaction if TYPE_CHECKING: from pyfragment.client import FragmentClient @@ -73,6 +73,7 @@ async def giveaway_stars( }, page_url=STARS_GIVEAWAY_PAGE, ) + required_payment_amount = parse_required_payment_amount(result, payment_method) req_id = result.get("req_id") if not req_id: raise FragmentAPIError(FragmentAPIError.NO_REQUEST_ID.format(context="Stars giveaway")) @@ -91,7 +92,12 @@ async def giveaway_stars( if transaction.get("need_verify"): raise VerificationError(VerificationError.KYC_REQUIRED) - tx_hash = await process_transaction(client, transaction) + tx_hash = await process_transaction( + client, + transaction, + payment_method=payment_method, + required_payment_amount=required_payment_amount, + ) return StarsGiveawayResult( transaction_id=tx_hash, channel=channel, diff --git a/pyfragment/methods/purchase_premium.py b/pyfragment/methods/purchase_premium.py index 6ba68d3..c3a4f11 100644 --- a/pyfragment/methods/purchase_premium.py +++ b/pyfragment/methods/purchase_premium.py @@ -14,7 +14,7 @@ from pyfragment.types import ( VerificationError, ) from pyfragment.types.constants import DEVICE, PREMIUM_PAGE, SUPPORTED_PAYMENT_METHODS, PaymentMethod -from pyfragment.utils import get_account_info, process_transaction +from pyfragment.utils import get_account_info, parse_required_payment_amount, process_transaction if TYPE_CHECKING: from pyfragment.client import FragmentClient @@ -71,6 +71,7 @@ async def purchase_premium( {"recipient": recipient, "months": months, "payment_method": payment_method}, page_url=PREMIUM_PAGE, ) + required_payment_amount = parse_required_payment_amount(result, payment_method) req_id = result.get("req_id") if not req_id: raise FragmentAPIError(FragmentAPIError.NO_REQUEST_ID.format(context="Premium purchase")) @@ -90,7 +91,12 @@ async def purchase_premium( if transaction.get("need_verify"): raise VerificationError(VerificationError.KYC_REQUIRED) - tx_hash = await process_transaction(client, transaction) + tx_hash = await process_transaction( + client, + transaction, + payment_method=payment_method, + required_payment_amount=required_payment_amount, + ) return PremiumResult(transaction_id=tx_hash, username=username, amount=months) except FragmentError: diff --git a/pyfragment/methods/purchase_stars.py b/pyfragment/methods/purchase_stars.py index 27dfc58..e51b0e8 100644 --- a/pyfragment/methods/purchase_stars.py +++ b/pyfragment/methods/purchase_stars.py @@ -14,7 +14,7 @@ from pyfragment.types import ( VerificationError, ) from pyfragment.types.constants import DEVICE, STARS_PAGE, SUPPORTED_PAYMENT_METHODS, PaymentMethod -from pyfragment.utils import get_account_info, process_transaction +from pyfragment.utils import get_account_info, parse_required_payment_amount, process_transaction if TYPE_CHECKING: from pyfragment.client import FragmentClient @@ -67,6 +67,7 @@ async def purchase_stars( {"recipient": recipient, "quantity": amount, "payment_method": payment_method}, page_url=STARS_PAGE, ) + required_payment_amount = parse_required_payment_amount(result, payment_method) req_id = result.get("req_id") if not req_id: raise FragmentAPIError(FragmentAPIError.NO_REQUEST_ID.format(context="Stars purchase")) @@ -86,7 +87,12 @@ async def purchase_stars( if transaction.get("need_verify"): raise VerificationError(VerificationError.KYC_REQUIRED) - tx_hash = await process_transaction(client, transaction) + tx_hash = await process_transaction( + client, + transaction, + payment_method=payment_method, + required_payment_amount=required_payment_amount, + ) return StarsResult(transaction_id=tx_hash, username=username, amount=amount) except FragmentError: diff --git a/pyfragment/types/constants.py b/pyfragment/types/constants.py index cd662d2..98191cc 100644 --- a/pyfragment/types/constants.py +++ b/pyfragment/types/constants.py @@ -16,8 +16,12 @@ SUPPORTED_WALLET_VERSIONS: frozenset[str] = frozenset(get_args(WalletVersion)) # Wallet class map — used to resolve the correct contract from WALLET_VERSION WALLET_CLASSES: dict[str, Any] = {"V4R2": WalletV4R2, "V5R1": WalletV5R1} -# Minimum wallet balance required to cover TON network gas fees. -MIN_TON_BALANCE: float = 0.056 +# Minimum minimum TON balance required for payments (transaction amount + gas reserve). +MIN_TON_BALANCE: float = 0.33 # 0.3 TON for transaction + 0.03 TON gas reserve + +# USDT (TON) jetton metadata used for payment-method balance checks. +USDT_TON_MASTER_ADDRESS: str = "EQCxE6mUtQJKFnGfaROTKOt1lZbDiiX1kCixRv7Nw2Id_sDs" +MIN_USDT_BALANCE: float = 0.75 # Default HTTP request timeout in seconds. DEFAULT_TIMEOUT: float = 30.0 diff --git a/pyfragment/types/exceptions.py b/pyfragment/types/exceptions.py index dba3c1e..6f8f621 100644 --- a/pyfragment/types/exceptions.py +++ b/pyfragment/types/exceptions.py @@ -131,11 +131,10 @@ class OperationError(FragmentError): class WalletError(OperationError): """Raised for TON wallet issues (connection, balance, account info).""" - LOW_BALANCE = ( - "Insufficient TON balance: {balance:.4f} TON available, {required:.4f} TON required " - "(transaction amount + {gas:.3f} TON gas reserve)." - ) - BALANCE_CHECK_FAILED = "Failed to fetch wallet balance: {exc}" + LOW_TON_BALANCE = "Insufficient TON balance: {balance:.4f} TON available, {required:.4f} TON required." + LOW_USDT_BALANCE = "Insufficient USDT balance: {balance:.4f} USDT available, {required:.4f} USDT required." + TON_BALANCE_CHECK_FAILED = "Failed to fetch TON balance: {exc}" + USDT_BALANCE_CHECK_FAILED = "Failed to fetch USDT balance: {exc}" ACCOUNT_INFO_FAILED = "Failed to retrieve wallet account info from TON network: {exc}" WALLET_INFO_FAILED = "Failed to retrieve wallet info from TON network: {exc}" diff --git a/pyfragment/types/results.py b/pyfragment/types/results.py index 1e4cecc..7ae9017 100644 --- a/pyfragment/types/results.py +++ b/pyfragment/types/results.py @@ -27,10 +27,14 @@ class WalletInfo: address: str state: str - balance: float + ton_balance: float + usdt_balance: float def __repr__(self) -> str: - return f"WalletInfo(address='{self.address}', state='{self.state}', balance={self.balance} TON)" + return ( + f"WalletInfo(address='{self.address}', state='{self.state}', " + f"ton_balance={self.ton_balance} TON, usdt_balance={self.usdt_balance} USDT)" + ) @dataclass diff --git a/pyfragment/utils/__init__.py b/pyfragment/utils/__init__.py index 45e96d7..cedd637 100644 --- a/pyfragment/utils/__init__.py +++ b/pyfragment/utils/__init__.py @@ -1,6 +1,6 @@ from pyfragment.utils.cookies import CookieResult, get_cookies_from_browser from pyfragment.utils.decoder import clean_decode -from pyfragment.utils.html import parse_auction_rows, parse_gift_items, parse_login_code +from pyfragment.utils.html import parse_auction_rows, parse_gift_items, parse_login_code, parse_required_payment_amount from pyfragment.utils.http import ( execute_transaction_request, fragment_request, @@ -17,6 +17,7 @@ __all__ = [ "parse_auction_rows", "parse_gift_items", "parse_login_code", + "parse_required_payment_amount", "execute_transaction_request", "fragment_request", "get_account_info", diff --git a/pyfragment/utils/decoder.py b/pyfragment/utils/decoder.py index d38a692..4285466 100644 --- a/pyfragment/utils/decoder.py +++ b/pyfragment/utils/decoder.py @@ -7,7 +7,7 @@ from ton_core import Cell from pyfragment.types import ParseError -def clean_decode(payload: str) -> str: +def clean_decode(payload: str) -> str | Cell: """Decode a base64-encoded BOC payload to a plain-text comment string. Fragment transaction payloads are BOC-serialised TVM cells. This function @@ -18,7 +18,8 @@ def clean_decode(payload: str) -> str: payload: Base64url-encoded BOC string (padding is added automatically). Returns: - Decoded comment string, or ``""`` for an empty payload. + Decoded comment string, ``""`` for an empty payload, or raw ``Cell`` + when payload is a non-UTF8 binary body. Raises: ParseError: If the payload cannot be decoded or parsed. @@ -28,10 +29,15 @@ def clean_decode(payload: str) -> str: return "" s += "=" * (-len(s) % 4) try: - boc = base64.b64decode(s) + # Fragment may return URL-safe base64 ("-"/"_") in transaction payloads. + boc = base64.b64decode(s, altchars=b"-_", validate=True) cell = Cell.one_from_boc(boc) sl = cell.begin_parse() - sl.load_uint(32) # op code — always 0 for text comment - return sl.load_snake_string().strip() + sl.load_uint(32) # op code + try: + return sl.load_snake_string().strip() + except UnicodeDecodeError: + # Some Fragment payloads are binary TVM cells rather than text comments. + return cell except Exception as exc: raise ParseError(ParseError.UNPARSEABLE.format(context="payload decode", exc=exc)) from exc diff --git a/pyfragment/utils/html.py b/pyfragment/utils/html.py index 28defd7..e968edf 100644 --- a/pyfragment/utils/html.py +++ b/pyfragment/utils/html.py @@ -161,3 +161,18 @@ def parse_gift_items(html: str) -> tuple[list[dict[str, Any]], int | None]: next_offset = int(next_offset_m.group(1)) if next_offset_m else None return items, next_offset + + +def parse_required_payment_amount(init_response: dict[str, Any], payment_method: str) -> float | None: + """Extract required payment amount from init*Request response. + + Fragment includes ``amount`` for both ``ton`` and ``usdt_ton`` init flows, + so this helper always reads ``amount``. + """ + del payment_method + + raw_amount = init_response.get("amount") + try: + return float(str(raw_amount)) + except (TypeError, ValueError): + return None diff --git a/pyfragment/utils/wallet.py b/pyfragment/utils/wallet.py index d5c4ee6..74c3f06 100644 --- a/pyfragment/utils/wallet.py +++ b/pyfragment/utils/wallet.py @@ -8,17 +8,99 @@ from typing import TYPE_CHECKING, Any from ton_core import NetworkGlobalID from tonutils.clients import TonapiClient +from tonutils.contracts.jetton import get_wallet_address_get_method, get_wallet_data_get_method from tonutils.exceptions import ProviderResponseError from pyfragment.types import TransactionError, WalletError, WalletInfo -from pyfragment.types.constants import MIN_TON_BALANCE, WALLET_CLASSES +from pyfragment.types.constants import ( + MIN_TON_BALANCE, + MIN_USDT_BALANCE, + USDT_TON_MASTER_ADDRESS, + WALLET_CLASSES, + PaymentMethod, +) from pyfragment.utils.decoder import clean_decode if TYPE_CHECKING: from pyfragment.client import FragmentClient -async def process_transaction(client: FragmentClient, transaction_data: dict[str, Any]) -> str: +async def _get_usdt_balance(ton: Any, wallet_address: str) -> float: + """Return wallet USDT balance via tonutils jetton get-methods.""" + try: + jetton_wallet_address = await get_wallet_address_get_method( + client=ton, + address=USDT_TON_MASTER_ADDRESS, + owner_address=wallet_address, + ) + wallet_data = await get_wallet_data_get_method(client=ton, address=jetton_wallet_address) + raw_balance = int(wallet_data[0]) if wallet_data else 0 + return float(raw_balance) / 1_000_000.0 + except ProviderResponseError as exc: + # No jetton wallet deployed yet -> effectively zero USDT balance. + if exc.code == 404: + return 0.0 + raise WalletError(WalletError.USDT_BALANCE_CHECK_FAILED.format(exc=exc)) from exc + except Exception as exc: + raise WalletError(WalletError.USDT_BALANCE_CHECK_FAILED.format(exc=exc)) from exc + + +async def _check_ton_payment_balance( + balance_ton: float, + amount_ton: float, + required_payment_amount: float | None, + ton: Any, + wallet_address: str, +) -> None: + """Validate balance requirements for TON payment method.""" + del ton + del wallet_address + + tx_price_ton = amount_ton + if required_payment_amount is not None and required_payment_amount > 0: + tx_price_ton = max(tx_price_ton, required_payment_amount) + + required_ton = max(tx_price_ton, MIN_TON_BALANCE) + if balance_ton < required_ton: + raise WalletError( + WalletError.LOW_TON_BALANCE.format( + balance=balance_ton, + required=required_ton, + ) + ) + + +async def _check_usdt_payment_balance( + balance_ton: float, + amount_ton: float, + required_payment_amount: float | None, + ton: Any, + wallet_address: str, +) -> None: + """Validate balance requirements for USDT payment method.""" + del amount_ton + + # USDT payment still needs TON for network fees. + if balance_ton < MIN_TON_BALANCE: + raise WalletError( + WalletError.LOW_TON_BALANCE.format( + balance=balance_ton, + required=MIN_TON_BALANCE, + ) + ) + + usdt_balance = await _get_usdt_balance(ton, wallet_address) + required_usdt = required_payment_amount if required_payment_amount is not None else MIN_USDT_BALANCE + if usdt_balance < required_usdt: + raise WalletError(WalletError.LOW_USDT_BALANCE.format(balance=usdt_balance, required=required_usdt)) + + +async def process_transaction( + client: FragmentClient, + transaction_data: dict[str, Any], + payment_method: PaymentMethod = "ton", + required_payment_amount: float | None = None, +) -> str: """Sign and broadcast a Fragment transaction to the TON network. Validates the payload structure, checks the wallet balance, decodes the @@ -27,6 +109,8 @@ async def process_transaction(client: FragmentClient, transaction_data: dict[str Args: client: Authenticated :class:`FragmentClient` instance. transaction_data: Raw transaction dict from ``execute_transaction_request``. + payment_method: Payment currency — ``"ton"`` or ``"usdt_ton"``. + required_payment_amount: Optional price from init*Request response. Returns: Normalised transaction hash string. @@ -45,20 +129,31 @@ async def process_transaction(client: FragmentClient, transaction_data: dict[str wallet_cls = WALLET_CLASSES[client.wallet_version] wallet, _, _, _ = wallet_cls.from_mnemonic(client=ton, mnemonic=client.seed) - # Check balance covers transaction amount + gas reserve + # Check balance covers selected payment flow requirements. try: await wallet.refresh() balance_ton = wallet.balance / 1_000_000_000 - required = amount_ton + MIN_TON_BALANCE - if balance_ton < required: - raise WalletError(WalletError.LOW_BALANCE.format(balance=balance_ton, required=required, gas=MIN_TON_BALANCE)) + wallet_address = wallet.address.to_str(False, False) + checkers = { + "ton": _check_ton_payment_balance, + "usdt_ton": _check_usdt_payment_balance, + } + checker = checkers[payment_method] + await checker( + balance_ton, + amount_ton, + required_payment_amount, + ton, + wallet_address, + ) except WalletError: raise except Exception as exc: - raise WalletError(WalletError.BALANCE_CHECK_FAILED.format(exc=exc)) from exc + raise WalletError(WalletError.TON_BALANCE_CHECK_FAILED.format(exc=exc)) from exc try: - payload = clean_decode(message["payload"]) + raw_payload = str(message.get("payload", "")) + payload = clean_decode(raw_payload) for attempt in range(3): try: @@ -130,7 +225,8 @@ async def get_wallet_info(client: FragmentClient) -> WalletInfo: client: Authenticated :class:`FragmentClient` instance. Returns: - :class:`WalletInfo` with ``address``, ``state``, and ``balance`` in TON. + :class:`WalletInfo` with ``address``, ``state``, ``balance`` in TON, + and ``usdt_balance`` in USDT. Raises: WalletError: If the wallet state cannot be fetched. @@ -140,10 +236,13 @@ async def get_wallet_info(client: FragmentClient) -> WalletInfo: wallet_cls = WALLET_CLASSES[client.wallet_version] wallet, _, _, _ = wallet_cls.from_mnemonic(client=ton, mnemonic=client.seed) await wallet.refresh() + wallet_address = wallet.address.to_str(False, False) + usdt_balance = await _get_usdt_balance(ton, wallet_address) return WalletInfo( address=wallet.address.to_str(is_user_friendly=True, is_bounceable=False), state=wallet.state.value, - balance=round(wallet.balance / 1_000_000_000, 4), + ton_balance=round(wallet.balance / 1_000_000_000, 4), + usdt_balance=round(usdt_balance, 4), ) except Exception as exc: raise WalletError(WalletError.WALLET_INFO_FAILED.format(exc=exc)) from exc diff --git a/tests/001_test_decode.py b/tests/001_test_decode.py index 8a60ac3..0541c4b 100644 --- a/tests/001_test_decode.py +++ b/tests/001_test_decode.py @@ -1,37 +1,62 @@ """Tests for clean_decode() — TON BOC payload decoding.""" +import base64 import re +from unittest.mock import patch import pytest +from ton_core import Cell from pyfragment.types import ParseError from pyfragment.utils.decoder import clean_decode -PAYLOADS = [ +PAYLOAD_CASES = [ pytest.param( "te6ccgEBAgEALwABTgAAAAAxMDAwMDAwIFRlbGVncmFtIFN0YXJzIAoKUmVmI1RQb01wegEABkM3ZQ", + True, id="stars", ), pytest.param( "te6ccgEBAgEANAABTgAAAABUZWxlZ3JhbSBQcmVtaXVtIGZvciAxIHllYXIgCgpSZWYjcgEAEE9OQnM2cmNt", + True, id="premium", ), pytest.param( "te6ccgEBAgEAMAABTgAAAABUZWxlZ3JhbSBhY2NvdW50IHRvcCB1cCAKClJlZiNrMXpDRQEACFkxd3g", + True, id="topup", ), + pytest.param( + "te6ccgEBAgEAfgABqA-KfqVP885dhccidjC3GwgBCkiH8LM_zUu0afyGCTWJwX1mDjdlf2rMa9UoQlD4UHUAF1jLlcMomlo5RJTwl8jnDDdfdhc7EgQQWPqFQ9IjyLPCAwEASgAAAAA1MCBUZWxlZ3JhbSBTdGFycyAKClJlZiNtOUpoWndBcFE", + False, + id="real_stars_50", + ), + pytest.param( + "te6ccgEBAgEANgABTgAAAABUZWxlZ3JhbSBQcmVtaXVtIGZvciAzIG1vbnRocyAKClJlZgEAFCMzcFdKdGJkYnU", + False, + id="real_premium_3m", + ), + pytest.param( + "te6ccgEBAwEAhgABqg-KfqWibdDaYaJCPUWWgvAIAQpIh_CzP81LtGn8hgk1icF9Zg43ZX9qzGvVKEJQ-FB1ABdYy5XDKJpaOUSU8JfI5ww3X3YXOxIEEFj6hUPSI8izwgMBAU4AAAAAMTAwMDAwIFRlbGVncmFtIFN0YXJzIAoKUmVmIzBoZ0RmNEYCAAQ5VA", + False, + id="real_stars_100k", + ), ] # Decode valid payload tests -@pytest.mark.parametrize("payload", PAYLOADS) -def test_decode_payload(payload: str) -> None: +@pytest.mark.parametrize(("payload", "strict_ref"), PAYLOAD_CASES) +def test_decode_payload(payload: str, strict_ref: bool) -> None: result = clean_decode(payload) - assert "Telegram" in result - assert re.search(r"Ref#[A-Za-z0-9]+", result), f"no Ref# in {result!r}" - assert all(ord(c) < 128 for c in result), f"non-ASCII chars in {result!r}" + if isinstance(result, str): + assert "Telegram" in result + if strict_ref: + assert re.search(r"Ref#[A-Za-z0-9]+", result), f"no Ref# in {result!r}" + assert all(ord(c) < 128 for c in result), f"non-ASCII chars in {result!r}" + else: + assert isinstance(result, Cell) # Edge case tests @@ -44,3 +69,64 @@ def test_empty_payload_returns_empty_string() -> None: def test_invalid_payload_raises_parse_error() -> None: with pytest.raises(ParseError): clean_decode("!!!not-valid-base64!!!") + + +def test_decode_payload_accepts_base64url_alphabet() -> None: + class _FakeSlice: + def load_uint(self, _: int) -> int: + return 0 + + def load_snake_string(self) -> str: + return "Telegram Stars Ref#abc" + + class _FakeCell: + def begin_parse(self) -> _FakeSlice: + return _FakeSlice() + + raw = b"\xfb\xef\xff\x00" + payload = base64.urlsafe_b64encode(raw).decode().rstrip("=") + + with patch("pyfragment.utils.decoder.Cell.one_from_boc", return_value=_FakeCell()) as mocked: + result = clean_decode(payload) + + mocked.assert_called_once_with(raw) + assert result == "Telegram Stars Ref#abc" + + +def test_clean_decode_returns_text_comment_when_utf8() -> None: + class _FakeSlice: + def load_uint(self, _: int) -> int: + return 0 + + def load_snake_string(self) -> str: + return "Telegram Premium Ref#abc" + + class _FakeCell: + def begin_parse(self) -> _FakeSlice: + return _FakeSlice() + + payload = base64.urlsafe_b64encode(b"\x00\x01").decode().rstrip("=") + with patch("pyfragment.utils.decoder.Cell.one_from_boc", return_value=_FakeCell()): + parsed = clean_decode(payload) + + assert parsed == "Telegram Premium Ref#abc" + + +def test_clean_decode_returns_cell_for_binary_payload() -> None: + class _FakeSlice: + def load_uint(self, _: int) -> int: + return 0 + + def load_snake_string(self) -> str: + raise UnicodeDecodeError("utf-8", b"\xff", 0, 1, "invalid start byte") + + class _FakeCell: + def begin_parse(self) -> _FakeSlice: + return _FakeSlice() + + payload = base64.urlsafe_b64encode(b"\x00\x01").decode().rstrip("=") + fake_cell: object = _FakeCell() + with patch("pyfragment.utils.decoder.Cell.one_from_boc", return_value=fake_cell): + parsed = clean_decode(payload) + + assert parsed is fake_cell diff --git a/tests/003_test_balance.py b/tests/003_test_balance.py index 6fa91b5..d3e0ea8 100644 --- a/tests/003_test_balance.py +++ b/tests/003_test_balance.py @@ -62,7 +62,7 @@ def _patch_wallet(wallet: MagicMock) -> Generator[None, None, None]: @pytest.mark.asyncio async def test_sufficient_balance_broadcasts() -> None: - wallet = _make_wallet(balance_nanotons=1_000_000_000) # 1 TON, needs 0.556 TON + wallet = _make_wallet(balance_nanotons=1_000_000_000) # 1 TON, above threshold with _patch_wallet(wallet), patch("pyfragment.utils.wallet.clean_decode", return_value="50 Telegram Stars"): result = await process_transaction(_make_client(), TRANSACTION_DATA) assert result == "abc123" @@ -71,7 +71,7 @@ async def test_sufficient_balance_broadcasts() -> None: @pytest.mark.asyncio async def test_insufficient_balance_raises() -> None: - wallet = _make_wallet(balance_nanotons=100_000_000) # 0.1 TON, needs 0.556 TON + wallet = _make_wallet(balance_nanotons=100_000_000) # 0.1 TON, below threshold with _patch_wallet(wallet): with pytest.raises(WalletError, match="required"): await process_transaction(_make_client(), TRANSACTION_DATA) @@ -80,7 +80,7 @@ async def test_insufficient_balance_raises() -> None: @pytest.mark.asyncio async def test_exact_minimum_balance_broadcasts() -> None: - wallet = _make_wallet(balance_nanotons=556_000_000) # exactly 0.5 + 0.056 TON + wallet = _make_wallet(balance_nanotons=500_000_000) # exactly transaction amount threshold with _patch_wallet(wallet), patch("pyfragment.utils.wallet.clean_decode", return_value="50 Telegram Stars"): result = await process_transaction(_make_client(), TRANSACTION_DATA) assert result == "abc123" @@ -88,7 +88,7 @@ async def test_exact_minimum_balance_broadcasts() -> None: @pytest.mark.asyncio async def test_one_nanoton_below_minimum_raises() -> None: - wallet = _make_wallet(balance_nanotons=555_999_999) # 1 nanoton below threshold + wallet = _make_wallet(balance_nanotons=499_999_999) # 1 nanoton below transaction amount threshold with _patch_wallet(wallet): with pytest.raises(WalletError, match="required"): await process_transaction(_make_client(), TRANSACTION_DATA) @@ -138,3 +138,41 @@ async def test_duplicate_seqno_raises_after_retries() -> None: with pytest.raises(TransactionError, match="seqno"): await process_transaction(_make_client(), TRANSACTION_DATA) assert wallet.transfer.call_count == 3 + + +@pytest.mark.asyncio +async def test_usdt_payment_requires_min_ton_gas_reserve() -> None: + wallet = _make_wallet(balance_nanotons=10_000_000) # 0.01 TON below MIN_TON_BALANCE + with _patch_wallet(wallet), patch("pyfragment.utils.wallet._get_usdt_balance", AsyncMock(return_value=100.0)): + with pytest.raises(WalletError, match="Insufficient TON balance"): + await process_transaction(_make_client(), TRANSACTION_DATA, payment_method="usdt_ton") + + +@pytest.mark.asyncio +async def test_usdt_payment_checks_usdt_balance() -> None: + wallet = _make_wallet(balance_nanotons=1_000_000_000) + transaction = { + "transaction": { + "messages": [ + { + "address": "0:852443f8599fe6a5da34fe43049ac4e0beb3071bb2bfb56635ea9421287c283a", + "amount": "50000000", + "payload": "", + } + ] + }, + "required_usdt": 12.5, + } + + with ( + _patch_wallet(wallet), + patch("pyfragment.utils.wallet.clean_decode", return_value=""), + patch("pyfragment.utils.wallet._get_usdt_balance", AsyncMock(return_value=5.0)), + ): + with pytest.raises(WalletError, match="Insufficient USDT balance"): + await process_transaction( + _make_client(), + transaction, + payment_method="usdt_ton", + required_payment_amount=12.5, + ) diff --git a/tests/004_test_stars.py b/tests/004_test_stars.py index ef99500..976ad4c 100644 --- a/tests/004_test_stars.py +++ b/tests/004_test_stars.py @@ -74,16 +74,19 @@ async def test_purchase_stars_passes_payment_method(client: FragmentClient) -> N FAKE_TRANSACTION, ] ) + proc_mock = AsyncMock(return_value=FAKE_TX_HASH) with ( patch.object(client, "call", call_mock), patch.object(_purchase_stars_mod, "get_account_info", AsyncMock(return_value=FAKE_ACCOUNT)), - patch.object(_purchase_stars_mod, "process_transaction", AsyncMock(return_value=FAKE_TX_HASH)), + patch.object(_purchase_stars_mod, "process_transaction", proc_mock), ): await client.purchase_stars("@user", amount=500, payment_method="usdt_ton") init_call = call_mock.await_args_list[2] assert init_call.args[0] == "initBuyStarsRequest" assert init_call.args[1]["payment_method"] == "usdt_ton" + assert proc_mock.await_args is not None + assert proc_mock.await_args.kwargs["payment_method"] == "usdt_ton" @pytest.mark.asyncio @@ -189,16 +192,19 @@ async def test_giveaway_stars_passes_payment_method(client: FragmentClient) -> N FAKE_TRANSACTION, ] ) + proc_mock = AsyncMock(return_value=FAKE_TX_HASH) with ( patch.object(client, "call", call_mock), patch.object(_giveaway_stars_mod, "get_account_info", AsyncMock(return_value=FAKE_ACCOUNT)), - patch.object(_giveaway_stars_mod, "process_transaction", AsyncMock(return_value=FAKE_TX_HASH)), + patch.object(_giveaway_stars_mod, "process_transaction", proc_mock), ): await client.giveaway_stars("@channel", winners=3, amount=1000, payment_method="usdt_ton") init_call = call_mock.await_args_list[1] assert init_call.args[0] == "initGiveawayStarsRequest" assert init_call.args[1]["payment_method"] == "usdt_ton" + assert proc_mock.await_args is not None + assert proc_mock.await_args.kwargs["payment_method"] == "usdt_ton" @pytest.mark.asyncio diff --git a/tests/005_test_premium.py b/tests/005_test_premium.py index 43df39e..3d111b1 100644 --- a/tests/005_test_premium.py +++ b/tests/005_test_premium.py @@ -71,16 +71,19 @@ async def test_purchase_premium_passes_payment_method(client: FragmentClient) -> FAKE_TRANSACTION, ] ) + proc_mock = AsyncMock(return_value=FAKE_TX_HASH) with ( patch.object(client, "call", call_mock), patch.object(_purchase_premium_mod, "get_account_info", AsyncMock(return_value=FAKE_ACCOUNT)), - patch.object(_purchase_premium_mod, "process_transaction", AsyncMock(return_value=FAKE_TX_HASH)), + patch.object(_purchase_premium_mod, "process_transaction", proc_mock), ): await client.purchase_premium("@user", months=6, payment_method="usdt_ton") init_call = call_mock.await_args_list[2] assert init_call.args[0] == "initGiftPremiumRequest" assert init_call.args[1]["payment_method"] == "usdt_ton" + assert proc_mock.await_args is not None + assert proc_mock.await_args.kwargs["payment_method"] == "usdt_ton" @pytest.mark.asyncio @@ -174,16 +177,19 @@ async def test_giveaway_premium_passes_payment_method(client: FragmentClient) -> FAKE_TRANSACTION, ] ) + proc_mock = AsyncMock(return_value=FAKE_TX_HASH) with ( patch.object(client, "call", call_mock), patch.object(_giveaway_premium_mod, "get_account_info", AsyncMock(return_value=FAKE_ACCOUNT)), - patch.object(_giveaway_premium_mod, "process_transaction", AsyncMock(return_value=FAKE_TX_HASH)), + patch.object(_giveaway_premium_mod, "process_transaction", proc_mock), ): await client.giveaway_premium("@channel", winners=10, months=6, payment_method="usdt_ton") init_call = call_mock.await_args_list[1] assert init_call.args[0] == "initGiveawayPremiumRequest" assert init_call.args[1]["payment_method"] == "usdt_ton" + assert proc_mock.await_args is not None + assert proc_mock.await_args.kwargs["payment_method"] == "usdt_ton" @pytest.mark.asyncio diff --git a/tests/007_test_wallet.py b/tests/007_test_wallet.py index d13c081..fbfb403 100644 --- a/tests/007_test_wallet.py +++ b/tests/007_test_wallet.py @@ -1,4 +1,4 @@ -"""Unit tests for get_wallet() — wallet address and TON balance lookup.""" +"""Unit tests for get_wallet() — wallet address/state with separate TON and USDT balances.""" from unittest.mock import AsyncMock, MagicMock, patch @@ -7,7 +7,7 @@ import pytest from pyfragment import FragmentClient, WalletInfo from tests.shared import FAKE_ADDRESS, FAKE_BALANCE_NANOTON -# Wallet mocked tests +# Wallet mocked tests (TON and USDT balances are returned separately) @pytest.mark.asyncio @@ -21,6 +21,7 @@ async def test_get_wallet_returns_wallet_info(client: FragmentClient) -> None: with ( patch("pyfragment.utils.wallet.TonapiClient") as mock_tonapi, patch("pyfragment.utils.wallet.WALLET_CLASSES") as mock_classes, + patch("pyfragment.utils.wallet._get_usdt_balance", AsyncMock(return_value=12.3456)), ): mock_tonapi.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) mock_tonapi.return_value.__aexit__ = AsyncMock(return_value=False) @@ -31,7 +32,8 @@ async def test_get_wallet_returns_wallet_info(client: FragmentClient) -> None: assert isinstance(result, WalletInfo) assert result.address == FAKE_ADDRESS assert result.state == "active" - assert result.balance == round(FAKE_BALANCE_NANOTON / 1_000_000_000, 4) + assert result.ton_balance == round(FAKE_BALANCE_NANOTON / 1_000_000_000, 4) + assert result.usdt_balance == 12.3456 @pytest.mark.asyncio @@ -45,6 +47,7 @@ async def test_get_wallet_balance_is_zero(client: FragmentClient) -> None: with ( patch("pyfragment.utils.wallet.TonapiClient") as mock_tonapi, patch("pyfragment.utils.wallet.WALLET_CLASSES") as mock_classes, + patch("pyfragment.utils.wallet._get_usdt_balance", AsyncMock(return_value=0.0)), ): mock_tonapi.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) mock_tonapi.return_value.__aexit__ = AsyncMock(return_value=False) @@ -52,5 +55,6 @@ async def test_get_wallet_balance_is_zero(client: FragmentClient) -> None: result = await client.get_wallet() - assert result.balance == 0.0 + assert result.ton_balance == 0.0 + assert result.usdt_balance == 0.0 assert result.state == "uninit" diff --git a/tests/015_test_payment_amount.py b/tests/015_test_payment_amount.py new file mode 100644 index 0000000..a202067 --- /dev/null +++ b/tests/015_test_payment_amount.py @@ -0,0 +1,21 @@ +"""Unit tests for init payment amount parsing.""" + +from pyfragment.utils.html import parse_required_payment_amount + + +def test_parse_required_payment_amount_ton_uses_amount() -> None: + init_response = {"amount": "0.326"} + assert parse_required_payment_amount(init_response, "ton") == 0.326 + + +def test_parse_required_payment_amount_usdt_uses_amount() -> None: + init_response = { + "amount": "0.00075", + "content": '0.75', + } + assert parse_required_payment_amount(init_response, "usdt_ton") == 0.00075 + + +def test_parse_required_payment_amount_usdt_falls_back_to_amount() -> None: + init_response = {"amount": "1.25", "content": "

no usd icon

"} + assert parse_required_payment_amount(init_response, "usdt_ton") == 1.25