feat: add wallet management features and enhance error handling; update tests and examples

.github/workflows/lint.yml

@@ -2,9 +2,9 @@ name: Lint
 
 on:
   push:
-    branches: [master]
+    branches: ["**"]
   pull_request:
-    branches: [master]
+    branches: ["**"]
 
 jobs:
   lint:

.github/workflows/publish.yml

@@ -2,7 +2,9 @@ name: Publish to PyPI
 
 on:
   push:
-    branches: [master]
+    branches: ["**"]
+  pull_request:
+    branches: ["**"]
 
 jobs:
   build:

.gitignore

@@ -23,7 +23,6 @@ logs/
 # System files
 .DS_Store
 Thumbs.db
-cookies.json
 
 # Testing & tooling artifacts
 .hypothesis/

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 | —                         |
 
 ---
 

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__":

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",

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)

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",
 ]

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):

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

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:
+            required = amount_ton + MIN_TON_BALANCE
-                raise WalletError(WalletError.LOW_BALANCE.format(balance=balance_ton))
+            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

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

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": {}})