Fix incorrect message serialization in transaction signing

Added jup_python_sdk/tests/ directory with tests for signing, serialization roundtrip, and key loading.

Author: shubhamessier

jup_python_sdk/clients/jupiter_client.py

@@ -5,6 +5,7 @@
 
 import base58
 import httpx
+from solders.message import to_bytes_versioned
 from solders.solders import Keypair, VersionedTransaction
 
 
@@ -93,7 +94,7 @@ def _sign_versioned_transaction(
 
         signers = list(versioned_transaction.signatures)
 
-        message_bytes = bytes(versioned_transaction.message)
+        message_bytes = to_bytes_versioned(versioned_transaction.message)
         your_signature = wallet.sign_message(message_bytes)
         signers[wallet_index] = your_signature
 

jup_python_sdk/tests/__init__.py

@@ -0,0 +1,161 @@
+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())