test_jupiter_client.py

import base64
import os
from unittest.mock import patch

import base58
import pytest
from solders.hash import Hash
from solders.instruction import AccountMeta, Instruction
from solders.keypair import Keypair
from solders.message import MessageV0, to_bytes_versioned
from solders.signature import Signature
from solders.system_program import TransferParams, transfer
from solders.transaction import VersionedTransaction

from jup_python_sdk.clients.jupiter_client import JupiterClient


@pytest.fixture
def wallet():
    return Keypair()


@pytest.fixture
def client(wallet):
    env_var = "TEST_PRIVATE_KEY"
    os.environ[env_var] = base58.b58encode(bytes(wallet)).decode()
    c = JupiterClient(api_key=None, private_key_env_var=env_var, timeout=10)
    yield c
    c.close()
    del os.environ[env_var]


def _build_versioned_tx(payer: Keypair, num_signers: int = 1):
    """Build a VersionedTransaction with the given number of signer slots."""
    extra_signers = [Keypair() for _ in range(num_signers - 1)]
    all_signers = [payer] + extra_signers

    # Build an instruction that requires all signers
    program_id = Keypair().pubkey()
    accounts = [
        AccountMeta(pubkey=s.pubkey(), is_signer=True, is_writable=True)
        for s in all_signers
    ]
    ix = Instruction(program_id, bytes(), accounts)

    blockhash = Hash.default()
    msg = MessageV0.try_compile(
        payer=payer.pubkey(),
        instructions=[ix],
        address_lookup_table_accounts=[],
        recent_blockhash=blockhash,
    )

    # Create a fully signed tx then zero out signatures to simulate
    # an unsigned transaction from the API
    signed_tx = VersionedTransaction(msg, all_signers)
    signatures = list(signed_tx.signatures)
    zeroed = [Signature.default()] * len(signatures)
    return VersionedTransaction.populate(msg, zeroed)


class TestSignVersionedTransaction:
    def test_single_signer(self, client, wallet):
        """Signing works for a transaction with one signature slot."""
        tx = _build_versioned_tx(wallet, num_signers=1)

        signed = client._sign_versioned_transaction(tx)

        assert len(signed.signatures) == 1
        assert signed.signatures[0] != Signature.default()

    def test_multi_signer_no_type_error(self, client, wallet):
        """Signing works for transactions with multiple signature slots (issue #3)."""
        tx = _build_versioned_tx(wallet, num_signers=3)
        assert len(tx.signatures) == 3

        signed = client._sign_versioned_transaction(tx)

        assert len(signed.signatures) == 3
        # Wallet signature should be filled in
        wallet_index = list(tx.message.account_keys).index(wallet.pubkey())
        assert signed.signatures[wallet_index] != Signature.default()
        # Other slots remain default (unsigned)
        for i, sig in enumerate(signed.signatures):
            if i != wallet_index:
                assert sig == Signature.default()

    def test_signature_is_valid(self, client, wallet):
        """The produced signature actually verifies against the message."""
        tx = _build_versioned_tx(wallet, num_signers=1)

        signed = client._sign_versioned_transaction(tx)

        msg_bytes = to_bytes_versioned(signed.message)
        sig = signed.signatures[0]
        # Verify by re-signing and comparing
        expected_sig = wallet.sign_message(msg_bytes)
        assert sig == expected_sig

    def test_uses_to_bytes_versioned(self, client, wallet):
        """Signing uses to_bytes_versioned, not bytes(), for correct message serialization."""
        tx = _build_versioned_tx(wallet, num_signers=1)

        with patch(
            "jup_python_sdk.clients.jupiter_client.to_bytes_versioned",
            wraps=to_bytes_versioned,
        ) as mock_fn:
            client._sign_versioned_transaction(tx)
            mock_fn.assert_called_once()

    def test_roundtrip_serialize_deserialize(self, client, wallet):
        """A signed transaction can be serialized to base64 and back."""
        tx = _build_versioned_tx(wallet, num_signers=2)

        signed = client._sign_versioned_transaction(tx)
        b64 = client._serialize_versioned_transaction(signed)

        recovered = VersionedTransaction.from_bytes(base64.b64decode(b64))
        assert recovered.signatures == signed.signatures


class TestSignBase64Transaction:
    def test_sign_base64_transaction(self, client, wallet):
        """_sign_base64_transaction decodes, signs, and returns a VersionedTransaction."""
        tx = _build_versioned_tx(wallet, num_signers=1)
        b64 = base64.b64encode(bytes(tx)).decode()

        signed = client._sign_base64_transaction(b64)

        assert isinstance(signed, VersionedTransaction)
        assert signed.signatures[0] != Signature.default()


class TestLoadPrivateKey:
    def test_base58_key(self, wallet):
        env_var = "TEST_PK_B58"
        os.environ[env_var] = base58.b58encode(bytes(wallet)).decode()
        c = JupiterClient(api_key=None, private_key_env_var=env_var, timeout=10)
        assert c._load_private_key_bytes() == bytes(wallet)
        del os.environ[env_var]

    def test_uint8_array_key(self, wallet):
        env_var = "TEST_PK_ARR"
        arr = list(bytes(wallet))
        os.environ[env_var] = str(arr)
        c = JupiterClient(api_key=None, private_key_env_var=env_var, timeout=10)
        assert c._load_private_key_bytes() == bytes(wallet)
        del os.environ[env_var]

    def test_invalid_key_raises(self):
        env_var = "TEST_PK_BAD"
        os.environ[env_var] = "not-a-valid-key!!!"
        c = JupiterClient(api_key=None, private_key_env_var=env_var, timeout=10)
        with pytest.raises(ValueError):
            c._load_private_key_bytes()
        del os.environ[env_var]


class TestGetPublicKey:
    def test_returns_correct_pubkey(self, client, wallet):
        assert client._get_public_key() == str(wallet.pubkey())