Support VersionedMessage in get_fee_for_message

CHANGELOG.md

@@ -4,7 +4,8 @@
 
 ## Added
 
-- Add VersionedTransaction support to `send_transaction` and `simulate_transaction` methods [(#334)](https://github.com/michaelhly/solana-py/pull/334):
+- Add VersionedTransaction support to `send_transaction` and `simulate_transaction` methods [(#334)](https://github.com/michaelhly/solana-py/pull/334)
+- Support VersionedMessage in `get_fee_for_message` methods [(#336)](https://github.com/michaelhly/solana-py/pull/336)
 
 ## Changed
 

src/solana/rpc/api.py

@@ -5,7 +5,7 @@
 from typing import Dict, List, Optional, Sequence, Union
 
 from solders.keypair import Keypair
-from solders.message import Message
+from solders.message import VersionedMessage
 from solders.pubkey import Pubkey
 from solders.rpc.responses import (
     GetAccountInfoMaybeJsonParsedResp,
@@ -411,7 +411,9 @@ def get_epoch_schedule(self) -> GetEpochScheduleResp:
         """
         return self._provider.make_request(self._get_epoch_schedule, GetEpochScheduleResp)
 
-    def get_fee_for_message(self, message: Message, commitment: Optional[Commitment] = None) -> GetFeeForMessageResp:
+    def get_fee_for_message(
+        self, message: VersionedMessage, commitment: Optional[Commitment] = None
+    ) -> GetFeeForMessageResp:
         """Returns the fee for a message.
 
         Args:

src/solana/rpc/async_api.py

@@ -4,7 +4,7 @@
 from typing import Dict, List, Optional, Sequence, Union
 
 from solders.keypair import Keypair
-from solders.message import Message
+from solders.message import VersionedMessage
 from solders.pubkey import Pubkey
 from solders.rpc.responses import (
     GetAccountInfoMaybeJsonParsedResp,
@@ -424,7 +424,7 @@ async def get_epoch_schedule(self) -> GetEpochScheduleResp:
         return await self._provider.make_request(self._get_epoch_schedule, GetEpochScheduleResp)
 
     async def get_fee_for_message(
-        self, message: Message, commitment: Optional[Commitment] = None
+        self, message: VersionedMessage, commitment: Optional[Commitment] = None
     ) -> GetFeeForMessageResp:
         """Returns the fee for a message.
 

src/solana/rpc/core.py

@@ -9,7 +9,7 @@
 
 from solders.account_decoder import UiAccountEncoding, UiDataSliceConfig
 from solders.commitment_config import CommitmentLevel
-from solders.message import Message
+from solders.message import MessageV0, VersionedMessage
 from solders.pubkey import Pubkey
 from solders.rpc.config import (
     RpcAccountInfoConfig,
@@ -254,8 +254,13 @@ def _get_epoch_info_body(self, commitment: Optional[Commitment]) -> GetEpochInfo
         config = RpcContextConfig(commitment=commitment_to_use)
         return GetEpochInfo(config)
 
-    def _get_fee_for_message_body(self, message: Message, commitment: Optional[Commitment]) -> GetFeeForMessage:
+    def _get_fee_for_message_body(
+        self, message: VersionedMessage, commitment: Optional[Commitment]
+    ) -> GetFeeForMessage:
         commitment_to_use = _COMMITMENT_TO_SOLDERS[commitment or self._commitment]
+        # weird mypy hack:
+        if isinstance(message, MessageV0):
+            return GetFeeForMessage(message, commitment_to_use)
         return GetFeeForMessage(message, commitment_to_use)
 
     def _get_inflation_governor_body(self, commitment: Optional[Commitment]) -> GetInflationGovernor:

tests/docker-compose.yml

@@ -1,7 +1,7 @@
 version: '3'
 services:
   localnet:
-    image: "solanalabs/solana:v1.10.25"
+    image: "solanalabs/solana:v1.13.6"
     ports:
       - "8899:8899"
       - "8900:8900"

tests/integration/test_async_http_client.py

@@ -376,6 +376,31 @@ async def test_get_fee_for_transaction_message(stubbed_sender, stubbed_receiver,
     assert fee_resp.value is not None
 
 
+@pytest.mark.integration
+async def test_get_fee_for_versioned_message(
+    stubbed_sender: Keypair, stubbed_receiver: Pubkey, test_http_client_async: AsyncClient
+):
+    """Test that gets a fee for a transaction using get_fee_for_message."""
+    # Get a recent blockhash
+    resp = await test_http_client_async.get_latest_blockhash()
+    recent_blockhash = resp.value.blockhash
+    assert recent_blockhash is not None
+    msg = MessageV0.try_compile(
+        payer=stubbed_sender.pubkey(),
+        instructions=[
+            sp.transfer(
+                sp.TransferParams(from_pubkey=stubbed_sender.pubkey(), to_pubkey=stubbed_receiver, lamports=1000)
+            )
+        ],
+        address_lookup_table_accounts=[],
+        recent_blockhash=recent_blockhash,
+    )
+    # get fee for transaction
+    fee_resp = await test_http_client_async.get_fee_for_message(msg)
+    assert_valid_response(fee_resp)
+    assert fee_resp.value is not None
+
+
 @pytest.mark.integration
 async def test_get_block_commitment(test_http_client_async):
     """Test get block commitment."""

tests/integration/test_http_client.py

@@ -347,6 +347,29 @@ def test_get_fee_for_transaction(stubbed_sender, stubbed_receiver, test_http_cli
     assert fee_resp.value is not None
 
 
+@pytest.mark.integration
+def test_get_fee_for_versioned_message(stubbed_sender: Keypair, stubbed_receiver: Pubkey, test_http_client: Client):
+    """Test that gets a fee for a transaction using get_fee_for_message."""
+    # Get a recent blockhash
+    resp = test_http_client.get_latest_blockhash()
+    recent_blockhash = resp.value.blockhash
+    assert recent_blockhash is not None
+    msg = MessageV0.try_compile(
+        payer=stubbed_sender.pubkey(),
+        instructions=[
+            sp.transfer(
+                sp.TransferParams(from_pubkey=stubbed_sender.pubkey(), to_pubkey=stubbed_receiver, lamports=1000)
+            )
+        ],
+        address_lookup_table_accounts=[],
+        recent_blockhash=recent_blockhash,
+    )
+    # get fee for transaction
+    fee_resp = test_http_client.get_fee_for_message(msg)
+    assert_valid_response(fee_resp)
+    assert fee_resp.value is not None
+
+
 @pytest.mark.integration
 def test_get_block_commitment(test_http_client: Client):
     """Test get block commitment."""