fix flaky test

solana/blockhash.py

@@ -23,14 +23,18 @@ def __init__(self, ttl: int = 60) -> None:
         self.unused_blockhashes: TTLCache = TTLCache(maxsize=maxsize, ttl=ttl)
         self.used_blockhashes: TTLCache = TTLCache(maxsize=maxsize, ttl=ttl)
 
-    def set(self, blockhash: Blockhash, slot: int) -> None:
+    def set(self, blockhash: Blockhash, slot: int, used_immediately: bool = False) -> None:
         """Update the cache.
 
         :param blockhash: new Blockhash value.
-        :param slot: the slot which the blockhash came from
-
+        :param slot: the slot which the blockhash came from.
+        :param used_immediately: whether the client used the blockhash immediately after fetching it.
 
         """
+        if used_immediately:
+            if slot not in self.used_blockhashes:
+                self.used_blockhashes[slot] = blockhash
+            return
         if slot in self.used_blockhashes or slot in self.unused_blockhashes:
             return
         self.unused_blockhashes[slot] = blockhash

solana/rpc/api.py

@@ -1044,7 +1044,7 @@ def send_transaction(
                     recent_blockhash = self.blockhash_cache.get()
                 except ValueError:
                     blockhash_resp = self.get_recent_blockhash()
-                    recent_blockhash = self._process_blockhash_resp(blockhash_resp)
+                    recent_blockhash = self._process_blockhash_resp(blockhash_resp, used_immediately=True)
             else:
                 blockhash_resp = self.get_recent_blockhash()
                 recent_blockhash = self.parse_recent_blockhash(blockhash_resp)
@@ -1054,7 +1054,7 @@ def send_transaction(
         txn_resp = self.send_raw_transaction(txn.serialize(), opts=opts)
         if self.blockhash_cache:
             blockhash_resp = self.get_recent_blockhash()
-            self._process_blockhash_resp(blockhash_resp)
+            self._process_blockhash_resp(blockhash_resp, used_immediately=False)
         return txn_resp
 
     def simulate_transaction(

solana/rpc/async_api.py

@@ -1042,7 +1042,7 @@ async def send_transaction(
                     recent_blockhash = self.blockhash_cache.get()
                 except ValueError:
                     blockhash_resp = await self.get_recent_blockhash()
-                    recent_blockhash = self._process_blockhash_resp(blockhash_resp)
+                    recent_blockhash = self._process_blockhash_resp(blockhash_resp, used_immediately=True)
             else:
                 blockhash_resp = await self.get_recent_blockhash()
                 recent_blockhash = self.parse_recent_blockhash(blockhash_resp)
@@ -1052,7 +1052,7 @@ async def send_transaction(
         txn_resp = await self.send_raw_transaction(txn.serialize(), opts=opts)
         if self.blockhash_cache:
             blockhash_resp = await self.get_recent_blockhash()
-            self._process_blockhash_resp(blockhash_resp)
+            self._process_blockhash_resp(blockhash_resp, used_immediately=False)
         return txn_resp
 
     async def simulate_transaction(

solana/rpc/core.py

@@ -360,9 +360,9 @@ def parse_recent_blockhash(blockhash_resp: types.RPCResponse) -> Blockhash:
             raise RuntimeError("failed to get recent blockhash")
         return Blockhash(blockhash_resp["result"]["value"]["blockhash"])
 
-    def _process_blockhash_resp(self, blockhash_resp: types.RPCResponse) -> Blockhash:
+    def _process_blockhash_resp(self, blockhash_resp: types.RPCResponse, used_immediately: bool) -> Blockhash:
         recent_blockhash = self.parse_recent_blockhash(blockhash_resp)
         if self.blockhash_cache:
             slot = blockhash_resp["result"]["context"]["slot"]
-            self.blockhash_cache.set(recent_blockhash, slot)
+            self.blockhash_cache.set(recent_blockhash, slot, used_immediately=used_immediately)
         return recent_blockhash

tests/integration/test_async_http_client.py

@@ -172,9 +172,14 @@ async def test_send_transaction_cached_blockhash(
             )
         )
     )
+    assert len(test_http_client_async_cached_blockhash.blockhash_cache.unused_blockhashes) == 0
+    assert len(test_http_client_async_cached_blockhash.blockhash_cache.used_blockhashes) == 0
     resp = await test_http_client_async_cached_blockhash.send_transaction(
         transfer_tx, async_stubbed_sender_cached_blockhash
     )
+    # we could have got a new blockhash or not depending on network latency and luck
+    assert len(test_http_client_async_cached_blockhash.blockhash_cache.unused_blockhashes) in (0, 1)
+    assert len(test_http_client_async_cached_blockhash.blockhash_cache.used_blockhashes) == 1
     assert_valid_response(resp)
     # Confirm transaction
     resp = await aconfirm_transaction(test_http_client_async_cached_blockhash, resp["result"])
@@ -207,7 +212,6 @@ async def test_send_transaction_cached_blockhash(
     resp = await test_http_client_async_cached_blockhash.get_balance(async_stubbed_sender_cached_blockhash.public_key())
     assert_valid_response(resp)
     assert resp["result"]["value"] == 9999994000
-    assert len(test_http_client_async_cached_blockhash.blockhash_cache.unused_blockhashes) == 1
 
     # Second transaction
     transfer_tx = Transaction().add(
@@ -225,6 +229,9 @@ async def test_send_transaction_cached_blockhash(
     resp = await test_http_client_async_cached_blockhash.send_transaction(
         transfer_tx, async_stubbed_sender_cached_blockhash
     )
+    # we could have got a new blockhash or not depending on network latency and luck
+    assert len(test_http_client_async_cached_blockhash.blockhash_cache.unused_blockhashes) in (0, 1)
+    assert len(test_http_client_async_cached_blockhash.blockhash_cache.used_blockhashes) in (1, 2)
     assert_valid_response(resp)
     # Confirm transaction
     resp = await aconfirm_transaction(test_http_client_async_cached_blockhash, resp["result"])
@@ -249,8 +256,6 @@ async def test_send_transaction_cached_blockhash(
     resp = await test_http_client_async_cached_blockhash.get_balance(async_stubbed_sender_cached_blockhash.public_key())
     assert_valid_response(resp)
     assert resp["result"]["value"] == 9999987000
-    assert len(test_http_client_async_cached_blockhash.blockhash_cache.unused_blockhashes) == 1
-    assert len(test_http_client_async_cached_blockhash.blockhash_cache.used_blockhashes) == 1
 
 
 @pytest.mark.integration

tests/integration/test_http_client.py

@@ -159,7 +159,12 @@ def test_send_transaction_cached_blockhash(
             )
         )
     )
+    assert len(test_http_client_cached_blockhash.blockhash_cache.unused_blockhashes) == 0
+    assert len(test_http_client_cached_blockhash.blockhash_cache.used_blockhashes) == 0
     resp = test_http_client_cached_blockhash.send_transaction(transfer_tx, stubbed_sender_cached_blockhash)
+    # we could have got a new blockhash or not depending on network latency and luck
+    assert len(test_http_client_cached_blockhash.blockhash_cache.unused_blockhashes) in (0, 1)
+    assert len(test_http_client_cached_blockhash.blockhash_cache.used_blockhashes) == 1
     assert_valid_response(resp)
     # Confirm transaction
     resp = confirm_transaction(test_http_client_cached_blockhash, resp["result"])
@@ -192,7 +197,6 @@ def test_send_transaction_cached_blockhash(
     resp = test_http_client_cached_blockhash.get_balance(stubbed_sender_cached_blockhash.public_key())
     assert_valid_response(resp)
     assert resp["result"]["value"] == 9999994000
-    assert len(test_http_client_cached_blockhash.blockhash_cache.unused_blockhashes) == 1
 
     # Second transaction
     transfer_tx = Transaction().add(
@@ -208,6 +212,9 @@ def test_send_transaction_cached_blockhash(
     assert_valid_response(resp)
     assert resp["result"]["value"] == 954
     resp = test_http_client_cached_blockhash.send_transaction(transfer_tx, stubbed_sender_cached_blockhash)
+    # we could have got a new blockhash or not depending on network latency and luck
+    assert len(test_http_client_cached_blockhash.blockhash_cache.unused_blockhashes) in (0, 1)
+    assert len(test_http_client_cached_blockhash.blockhash_cache.used_blockhashes) in (1, 2)
     assert_valid_response(resp)
     # Confirm transaction
     resp = confirm_transaction(test_http_client_cached_blockhash, resp["result"])