From d8f0240807f9aab2f5c989aa6a84f4462d3f054b Mon Sep 17 00:00:00 2001 From: bohd4nx Date: Sun, 15 Mar 2026 22:41:52 +0200 Subject: [PATCH] feat: add wallet management features and enhance error handling; update tests and examples --- .github/workflows/lint.yml | 4 +- .github/workflows/publish.yml | 4 +- .gitignore | 1 - README.md | 1 + examples/{client_init.py => get_wallet.py} | 6 +- pyfragment/__init__.py | 7 ++ pyfragment/client.py | 16 +++ pyfragment/types/__init__.py | 3 +- pyfragment/types/exceptions.py | 4 +- pyfragment/types/results.py | 11 +- pyfragment/utils/wallet.py | 46 ++++++-- tests/002_test_client.py | 13 +++ tests/004_test_balance.py | 127 +++++++++++++++++++++ 13 files changed, 225 insertions(+), 18 deletions(-) rename examples/{client_init.py => get_wallet.py} (81%) create mode 100644 tests/004_test_balance.py diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index ed1ead3..2a60631 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -2,9 +2,9 @@ name: Lint on: push: - branches: [master] + branches: ["**"] pull_request: - branches: [master] + branches: ["**"] jobs: lint: diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index d8955a1..d0f7ae0 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -2,7 +2,9 @@ name: Publish to PyPI on: push: - branches: [master] + branches: ["**"] + pull_request: + branches: ["**"] jobs: build: diff --git a/.gitignore b/.gitignore index 082dca9..ca55ba5 100644 --- a/.gitignore +++ b/.gitignore @@ -23,7 +23,6 @@ logs/ # System files .DS_Store Thumbs.db -cookies.json # Testing & tooling artifacts .hypothesis/ diff --git a/README.md b/README.md index 2e920ed..3e0ed5f 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,7 @@ See the [`examples/`](examples/) folder for ready-to-run scripts. | `gift_premium(username, months, show_sender=True)` | `PremiumResult` | Gift Telegram Premium subscription | `months`: 3, 6, or 12 | | `gift_stars(username, amount, show_sender=True)` | `StarsResult` | Gift Telegram Stars | `amount`: 50–1,000,000 | | `topup_ton(username, amount, show_sender=True)` | `AdsTopupResult` | Top up Telegram Ads balance | `amount`: 1–1,000,000,000 | +| `get_wallet()` | `WalletInfo` | Get wallet address, state, balance | — | --- diff --git a/examples/client_init.py b/examples/get_wallet.py similarity index 81% rename from examples/client_init.py rename to examples/get_wallet.py index 8fea035..936636e 100644 --- a/examples/client_init.py +++ b/examples/get_wallet.py @@ -27,9 +27,13 @@ async def main() -> None: wallet_version="V5R1", # or "V4R2" ) + wallet = await client.get_wallet() + print("FragmentClient initialized") print(" %-16s %s" % ("Wallet version:", client.wallet_version)) - print(" %-16s %s..." % ("API key:", client.api_key[:8])) + print(" %-16s %s" % ("Address:", wallet.address)) + print(" %-16s %s" % ("State:", wallet.state)) + print(" %-16s %s TON" % ("Balance:", wallet.balance)) if __name__ == "__main__": diff --git a/pyfragment/__init__.py b/pyfragment/__init__.py index 31f8445..9495146 100644 --- a/pyfragment/__init__.py +++ b/pyfragment/__init__.py @@ -3,6 +3,8 @@ # This source code is licensed under the MIT License found in the # LICENSE file in the root directory of this source tree. +from importlib.metadata import version + from pyfragment.client import FragmentClient from pyfragment.types import ( AdsTopupResult, @@ -21,13 +23,18 @@ from pyfragment.types import ( UserNotFoundError, VerificationError, WalletError, + WalletInfo, ) +__version__: str = version("pyfragment") + __all__ = [ + "__version__", "FragmentClient", "AdsTopupResult", "PremiumResult", "StarsResult", + "WalletInfo", "ClientError", "ConfigurationError", "CookieError", diff --git a/pyfragment/client.py b/pyfragment/client.py index 9884e2e..677f90a 100644 --- a/pyfragment/client.py +++ b/pyfragment/client.py @@ -11,8 +11,10 @@ from pyfragment.types import ( CookieError, PremiumResult, StarsResult, + WalletInfo, WalletVersion, ) +from pyfragment.utils.wallet import get_wallet_info class FragmentClient: @@ -36,6 +38,7 @@ class FragmentClient: api_key="AAABBB...", cookies={"stel_ssid": "...", "stel_dt": "...", ...}, ) + print(await client.get_wallet()) result = await client.gift_premium("@username", months=6) print(result.transaction_id) """ @@ -51,6 +54,10 @@ class FragmentClient: if missing: raise ConfigurationError(ConfigurationError.MISSING_VARS.format(keys=", ".join(missing))) + word_count = len(seed.split()) + if word_count not in (12, 18, 24): + raise ConfigurationError(ConfigurationError.INVALID_MNEMONIC.format(count=word_count)) + if isinstance(cookies, str): try: cookies = json.loads(cookies) @@ -112,3 +119,12 @@ class FragmentClient: :class:`AdsTopupResult` with ``transaction_id``, ``username``, ``amount``, ``timestamp``. """ return await topup_ton(self, username, amount, show_sender) + + async def get_wallet(self) -> WalletInfo: + """Return the address, state and balance of the TON wallet. + + Returns: + :class:`WalletInfo` with ``address`` (``"UQ..."``), ``state`` + (``"active"``, ``"uninit"``, or ``"frozen"``), and ``balance`` in TON. + """ + return await get_wallet_info(self) diff --git a/pyfragment/types/__init__.py b/pyfragment/types/__init__.py index 5cb878e..c254997 100644 --- a/pyfragment/types/__init__.py +++ b/pyfragment/types/__init__.py @@ -25,7 +25,7 @@ from pyfragment.types.exceptions import ( VerificationError, WalletError, ) -from pyfragment.types.results import AdsTopupResult, PremiumResult, StarsResult +from pyfragment.types.results import AdsTopupResult, PremiumResult, StarsResult, WalletInfo __all__ = [ # constants @@ -58,4 +58,5 @@ __all__ = [ "AdsTopupResult", "PremiumResult", "StarsResult", + "WalletInfo", ] diff --git a/pyfragment/types/exceptions.py b/pyfragment/types/exceptions.py index 6a8de78..b25fdbe 100644 --- a/pyfragment/types/exceptions.py +++ b/pyfragment/types/exceptions.py @@ -11,6 +11,7 @@ class ConfigurationError(ClientError): MISSING_VARS = "Missing required parameter(s): {keys}." UNSUPPORTED_VERSION = "Unsupported wallet_version '{version}'. Must be one of: {supported}." + INVALID_MNEMONIC = "Invalid mnemonic: got {count} words, expected 12, 18, or 24." INVALID_MONTHS = "Invalid duration. Choose 3, 6, or 12 months." INVALID_STARS_AMOUNT = "Amount must be an integer between 50 and 1 000 000 stars." INVALID_TON_AMOUNT = "Amount must be an integer between 1 and 1 000 000 000 TON." @@ -78,9 +79,10 @@ class OperationError(FragmentError): class WalletError(OperationError): """Raised for TON wallet issues (connection, balance, account info).""" - LOW_BALANCE = "TON wallet balance is too low: {balance:.2f} TON. Minimum required is 0.056 TON." + LOW_BALANCE = "TON wallet balance is too low: {balance:.4f} TON available, {required:.4f} TON required." BALANCE_CHECK_FAILED = "Wallet balance check failed: {exc}" ACCOUNT_INFO_FAILED = "Failed to retrieve wallet account info: {exc}" + WALLET_INFO_FAILED = "Failed to retrieve wallet info: {exc}" class UnexpectedError(OperationError): diff --git a/pyfragment/types/results.py b/pyfragment/types/results.py index 234f7cc..55d4add 100644 --- a/pyfragment/types/results.py +++ b/pyfragment/types/results.py @@ -1,7 +1,16 @@ import time from dataclasses import dataclass, field -__all__ = ["AdsTopupResult", "PremiumResult", "StarsResult"] +__all__ = ["AdsTopupResult", "PremiumResult", "StarsResult", "WalletInfo"] + + +@dataclass +class WalletInfo: + """Wallet state returned by :meth:`FragmentClient.get_wallet`.""" + + address: str + state: str + balance: float @dataclass diff --git a/pyfragment/utils/wallet.py b/pyfragment/utils/wallet.py index 57ae832..3a8fded 100644 --- a/pyfragment/utils/wallet.py +++ b/pyfragment/utils/wallet.py @@ -5,16 +5,13 @@ from tonutils.clients import TonapiClient from tonutils.types import NetworkGlobalID from pyfragment.types import MIN_TON_BALANCE, WALLET_CLASSES, TransactionError, WalletError +from pyfragment.types.results import WalletInfo from pyfragment.utils.decoder import clean_decode if TYPE_CHECKING: from pyfragment.client import FragmentClient -def _init_ton_client(client: "FragmentClient") -> TonapiClient: - return TonapiClient(network=NetworkGlobalID.MAINNET, api_key=client.api_key) - - async def process_transaction(client: "FragmentClient", transaction_data: dict) -> str: """Sign and broadcast a Fragment transaction to the TON network. @@ -35,26 +32,29 @@ async def process_transaction(client: "FragmentClient", transaction_data: dict) if "transaction" not in transaction_data or "messages" not in transaction_data["transaction"]: raise TransactionError(TransactionError.INVALID_PAYLOAD) + message = transaction_data["transaction"]["messages"][0] + amount_ton = int(message["amount"]) / 1_000_000_000 + # TODO: Investigate 406 'inbound external message rejected before smart-contract execution'. # This happens when the previous transaction's seqno hasn't been confirmed on-chain yet, # causing the wallet contract to reject the new message. - async with _init_ton_client(client) as ton: + async with TonapiClient(network=NetworkGlobalID.MAINNET, api_key=client.api_key) as ton: wallet_cls = WALLET_CLASSES[client.wallet_version] wallet, _, _, _ = wallet_cls.from_mnemonic(client=ton, mnemonic=client.seed) - # Check balance before broadcasting + # Check balance covers transaction amount + gas reserve try: await wallet.refresh() balance_ton = wallet.balance / 1_000_000_000 - if balance_ton < MIN_TON_BALANCE: - raise WalletError(WalletError.LOW_BALANCE.format(balance=balance_ton)) + required = amount_ton + MIN_TON_BALANCE + if balance_ton < required: + raise WalletError(WalletError.LOW_BALANCE.format(balance=balance_ton, required=required)) except WalletError: raise except Exception as exc: raise WalletError(WalletError.BALANCE_CHECK_FAILED.format(exc=exc)) from exc try: - message = transaction_data["transaction"]["messages"][0] payload = clean_decode(message["payload"]) result = await wallet.transfer( @@ -85,7 +85,7 @@ async def get_account_info(client: "FragmentClient") -> dict[str, Any]: Raises: WalletError: If account info cannot be retrieved. """ - async with _init_ton_client(client) as ton: + async with TonapiClient(network=NetworkGlobalID.MAINNET, api_key=client.api_key) as ton: try: wallet_cls = WALLET_CLASSES[client.wallet_version] wallet, pub_key, _, _ = wallet_cls.from_mnemonic(client=ton, mnemonic=client.seed) @@ -98,3 +98,29 @@ async def get_account_info(client: "FragmentClient") -> dict[str, Any]: } except Exception as exc: raise WalletError(WalletError.ACCOUNT_INFO_FAILED.format(exc=exc)) from exc + + +async def get_wallet_info(client: "FragmentClient") -> "WalletInfo": + """Return the address, state and balance of the TON wallet. + + Args: + client: Authenticated :class:`FragmentClient` instance. + + Returns: + :class:`WalletInfo` with ``address``, ``state``, and ``balance`` in TON. + + Raises: + WalletError: If the wallet state cannot be fetched. + """ + async with TonapiClient(network=NetworkGlobalID.MAINNET, api_key=client.api_key) as ton: + try: + wallet_cls = WALLET_CLASSES[client.wallet_version] + wallet, _, _, _ = wallet_cls.from_mnemonic(client=ton, mnemonic=client.seed) + await wallet.refresh() + 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), + ) + except Exception as exc: + raise WalletError(WalletError.WALLET_INFO_FAILED.format(exc=exc)) from exc diff --git a/tests/002_test_client.py b/tests/002_test_client.py index d62f144..79b7406 100644 --- a/tests/002_test_client.py +++ b/tests/002_test_client.py @@ -79,3 +79,16 @@ def test_whitespace_cookie_value_raises() -> None: bad = {**VALID_COOKIES, "stel_ton_token": " "} with pytest.raises(CookieError): FragmentClient(seed=VALID_SEED, api_key=VALID_API_KEY, cookies=bad) + + +def test_invalid_mnemonic_length_raises() -> None: + bad_seed = " ".join(["word"] * 23) + with pytest.raises(ConfigurationError): + FragmentClient(seed=bad_seed, api_key=VALID_API_KEY, cookies=VALID_COOKIES) + + +def test_valid_mnemonic_lengths() -> None: + for length in (12, 18, 24): + seed = " ".join(["abandon"] * (length - 1) + ["about"]) + client = FragmentClient(seed=seed, api_key=VALID_API_KEY, cookies=VALID_COOKIES) + assert len(client.seed.split()) == length diff --git a/tests/004_test_balance.py b/tests/004_test_balance.py new file mode 100644 index 0000000..e2023ee --- /dev/null +++ b/tests/004_test_balance.py @@ -0,0 +1,127 @@ +"""Unit tests for process_transaction() — balance checks before broadcast.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pyfragment.types import WalletError +from pyfragment.utils.wallet import process_transaction + +VALID_SEED = "abandon " * 23 + "about" + +# Minimal transaction payload: 0.5 TON = 500_000_000 nanotons +TRANSACTION_DATA = { + "transaction": { + "messages": [ + { + "address": "0:852443f8599fe6a5da34fe43049ac4e0beb3071bb2bfb56635ea9421287c283a", + "amount": "500000000", + "payload": "", + } + ] + } +} + + +def _make_client(api_key: str = "test_key") -> MagicMock: + client = MagicMock() + client.api_key = api_key + client.seed = VALID_SEED.split() + client.wallet_version = "V5R1" + return client + + +def _make_wallet(balance_nanotons: int) -> MagicMock: + wallet = MagicMock() + wallet.refresh = AsyncMock() + wallet.balance = balance_nanotons + wallet.transfer = AsyncMock(return_value=MagicMock(normalized_hash="abc123")) + return wallet + + +@pytest.mark.asyncio +async def test_sufficient_balance_broadcasts() -> None: + # 0.5 TON amount + 0.056 TON gas = 0.556 TON required; wallet has 1 TON + client = _make_client() + wallet = _make_wallet(balance_nanotons=1_000_000_000) + + with ( + patch("pyfragment.utils.wallet.TonapiClient") as mock_tonapi, + patch("pyfragment.utils.wallet.WALLET_CLASSES") as mock_classes, + patch("pyfragment.utils.wallet.clean_decode", return_value="50 Telegram Stars"), + ): + mock_tonapi.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) + mock_tonapi.return_value.__aexit__ = AsyncMock(return_value=False) + mock_classes["V5R1"].from_mnemonic.return_value = (wallet, MagicMock(), None, None) + + result = await process_transaction(client, TRANSACTION_DATA) + + assert result == "abc123" + wallet.transfer.assert_called_once() + + +@pytest.mark.asyncio +async def test_insufficient_balance_raises_wallet_error() -> None: + # wallet has 0.1 TON, needs 0.556 TON + client = _make_client() + wallet = _make_wallet(balance_nanotons=100_000_000) + + with ( + patch("pyfragment.utils.wallet.TonapiClient") as mock_tonapi, + patch("pyfragment.utils.wallet.WALLET_CLASSES") as mock_classes, + ): + mock_tonapi.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) + mock_tonapi.return_value.__aexit__ = AsyncMock(return_value=False) + mock_classes["V5R1"].from_mnemonic.return_value = (wallet, MagicMock(), None, None) + + with pytest.raises(WalletError, match="required"): + await process_transaction(client, TRANSACTION_DATA) + + wallet.transfer.assert_not_called() + + +@pytest.mark.asyncio +async def test_exactly_minimum_balance_broadcasts() -> None: + # exactly amount + gas: 500_000_000 + 56_000_000 = 556_000_000 nanotons + client = _make_client() + wallet = _make_wallet(balance_nanotons=556_000_000) + + with ( + patch("pyfragment.utils.wallet.TonapiClient") as mock_tonapi, + patch("pyfragment.utils.wallet.WALLET_CLASSES") as mock_classes, + patch("pyfragment.utils.wallet.clean_decode", return_value="50 Telegram Stars"), + ): + mock_tonapi.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) + mock_tonapi.return_value.__aexit__ = AsyncMock(return_value=False) + mock_classes["V5R1"].from_mnemonic.return_value = (wallet, MagicMock(), None, None) + + result = await process_transaction(client, TRANSACTION_DATA) + + assert result == "abc123" + + +@pytest.mark.asyncio +async def test_one_nanoton_below_minimum_raises() -> None: + # 556_000_000 - 1 nanoton: just below threshold + client = _make_client() + wallet = _make_wallet(balance_nanotons=555_999_999) + + with ( + patch("pyfragment.utils.wallet.TonapiClient") as mock_tonapi, + patch("pyfragment.utils.wallet.WALLET_CLASSES") as mock_classes, + ): + mock_tonapi.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) + mock_tonapi.return_value.__aexit__ = AsyncMock(return_value=False) + mock_classes["V5R1"].from_mnemonic.return_value = (wallet, MagicMock(), None, None) + + with pytest.raises(WalletError, match="required"): + await process_transaction(client, TRANSACTION_DATA) + + +@pytest.mark.asyncio +async def test_invalid_payload_raises_transaction_error() -> None: + from pyfragment.types import TransactionError + + client = _make_client() + with pytest.raises(TransactionError): + await process_transaction(client, {"transaction": {}})