fix: forked timestamp

ape_titanoboa/provider.py

@@ -23,6 +23,7 @@
 
 if TYPE_CHECKING:
     from ape.api import ForkedNetworkAPI, ReceiptAPI, TestAccountAPI, TraceAPI, TransactionAPI
+    from ape.api.networks import ProviderContextManager
     from ape.types import AddressType, BlockID, ContractCode, ContractLog, LogFilter, SnapshotID
     from ape_test.accounts import TestAccount
     from ape_test.config import ApeTestConfig
@@ -449,23 +450,29 @@ def fork_config(self) -> "BoaForkConfig":
     def forked_network(self) -> "ForkedNetworkAPI":
         return self.network  # type: ignore
 
-    @cached_property
-    def fork_url(self) -> str:
+    @property
+    def _forked_connection(self) -> "ProviderContextManager":
         if provider := self.fork_config.upstream_provider:
-            ctx = self.forked_network.upstream_network.use_provider(provider)
+            return self.forked_network.upstream_network.use_provider(provider)
         else:
-            ctx = self.forked_network.upstream_network.use_default_provider()
+            return self.forked_network.upstream_network.use_default_provider()
 
-        with ctx as upstream_provider:
+    @cached_property
+    def fork_url(self) -> str:
+        with self._forked_connection as upstream_provider:
             return upstream_provider.http_uri
 
     @property
     def block_identifier(self) -> Optional["BlockID"]:
         return self.fork_config.get("block_identifier")
 
+    @cached_property
+    def forked_block_start(self) -> "BlockAPI":
+        with self._forked_connection as upstream_provider:
+            return upstream_provider.get_block(self.block_identifier or "latest")
+
     def connect(self):
         self.fork.__enter__()
-        super().connect()
 
     def disconnect(self):
         self.fork.__exit__()
@@ -474,6 +481,22 @@ def disconnect(self):
         for cached_prop in ("fork", "fork_url", "fork_config"):
             self.__dict__.pop(cached_prop, None)
 
+    def get_block(self, block_id: "BlockID") -> BlockAPI:
+        start_number = self.forked_block_start.number or 0
+        try:
+            result = super().get_block(block_id)
+        except Exception:
+            # Try upstream.
+            with self._forked_connection as upstream_provider:
+                return upstream_provider.get_block(block_id)
+
+        if result.number is not None and result.number <= start_number:
+            # Correct data such as timestamp, which may be necessary for some tests.
+            with self._forked_connection as upstream_provider:
+                return upstream_provider.get_block(result.number)
+
+        return result
+
     def make_request(self, rpc: str, parameters: Optional[Iterable] = None) -> Any:
         # TODO: Make request directly to upstream URL.
         pass

tests/test_fork_provider.py

@@ -15,3 +15,7 @@ def test_get_latest_block(chain):
     expected = 7341111  # From configuration.
     actual = chain.blocks.height
     assert actual == expected
+
+    # NOTE: It is very important the timestamp is accurate
+    #   for any on-chain state that requires it.
+    assert chain.blocks.head.timestamp == 1734992352