diff --git a/ouroboros-consensus-test/test-storage/Test/Ouroboros/Storage/ChainDB/StateMachine.hs b/ouroboros-consensus-test/test-storage/Test/Ouroboros/Storage/ChainDB/StateMachine.hs
index 93d74b9ff8f..62111b7ca9d 100644
--- a/ouroboros-consensus-test/test-storage/Test/Ouroboros/Storage/ChainDB/StateMachine.hs
+++ b/ouroboros-consensus-test/test-storage/Test/Ouroboros/Storage/ChainDB/StateMachine.hs
@@ -35,7 +35,7 @@ import           Data.Functor.Identity (Identity)
 import           Data.List (sortOn)
 import qualified Data.List.NonEmpty as NE
 import qualified Data.Map as Map
-import           Data.Maybe (fromJust, fromMaybe)
+import           Data.Maybe (fromMaybe)
 import           Data.Ord (Down (..))
 import           Data.Proxy
 import           Data.Sequence.Strict (StrictSeq)
@@ -380,7 +380,10 @@ run env@ChainDBEnv { varDB, .. } cmd =
     advanceAndAdd ChainDBState { chainDB } newCurSlot blk = do
       atomically $ modifyTVar varCurSlot (max newCurSlot)
       -- `blockProcessed` always returns 'Just'
-      fromJust <$> addBlock chainDB InvalidBlockPunishment.noPunishment blk
+      res <- addBlock chainDB InvalidBlockPunishment.noPunishment blk
+      return $ case res of
+        Nothing -> error "advanceAndAdd: block not added"
+        Just pt -> pt
 
     wipeVolatileDB :: ChainDBState m blk -> m (Point blk)
     wipeVolatileDB st = do
diff --git a/ouroboros-consensus/src/Ouroboros/Consensus/Storage/ChainDB/API.hs b/ouroboros-consensus/src/Ouroboros/Consensus/Storage/ChainDB/API.hs
index 1ed4a5d5ba2..800c75ff7b1 100644
--- a/ouroboros-consensus/src/Ouroboros/Consensus/Storage/ChainDB/API.hs
+++ b/ouroboros-consensus/src/Ouroboros/Consensus/Storage/ChainDB/API.hs
@@ -439,6 +439,8 @@ addBlockWaitWrittenToDisk chainDB punish blk = do
 
 -- | Add a block synchronously: wait until the block has been processed (see
 -- 'blockProcessed'). The new tip of the ChainDB is returned.
+--
+-- Note: this is a partial function, only to support tests.
 addBlock :: IOLike m => ChainDB m blk -> InvalidBlockPunishment m -> blk -> m (Maybe (Point blk))
 addBlock chainDB punish blk = do
     promise <- addBlockAsync chainDB punish blk
@@ -446,6 +448,8 @@ addBlock chainDB punish blk = do
 
 -- | Add a block synchronously. Variant of 'addBlock' that doesn't return the
 -- new tip of the ChainDB.
+--
+-- Note: this is a partial function, only to support tests.
 addBlock_ :: IOLike m => ChainDB m blk -> InvalidBlockPunishment m -> blk -> m ()
 addBlock_  = void ..: addBlock