Share partial accessor functions used in tests

ouroboros-consensus-diffusion/ouroboros-consensus-diffusion.cabal

@@ -266,6 +266,7 @@ test-suite consensus-test
     Test.Consensus.PointSchedule.SinglePeer
     Test.Consensus.PointSchedule.SinglePeer.Indices
     Test.Consensus.PointSchedule.Tests
+    Test.Util.PartialAccessors
     Test.Util.TersePrinting
 
   build-depends:

ouroboros-consensus-diffusion/test/consensus-test/Test/Consensus/Genesis/Tests/DensityDisconnect.hs

@@ -58,6 +58,7 @@ import           Test.QuickCheck.Extras (unsafeMapSuchThatJust)
 import           Test.Tasty
 import           Test.Tasty.QuickCheck
 import           Test.Util.Orphans.IOLike ()
+import           Test.Util.PartialAccessors
 import           Test.Util.TersePrinting (terseHFragment, terseHeader)
 import           Test.Util.TestBlock (TestBlock)
 import           Test.Util.TestEnv (adjustQuickCheckMaxSize,
@@ -481,16 +482,6 @@ prop_densityDisconnectTriggersChainSel =
     )
 
   where
-    getOnlyBranch :: BlockTree blk -> BlockTreeBranch blk
-    getOnlyBranch BlockTree {btBranches} = case btBranches of
-      [branch] -> branch
-      _        -> error "tree must have exactly one alternate branch"
-
-    getTrunkTip :: HasHeader blk => BlockTree blk -> blk
-    getTrunkTip tree = case btTrunk tree of
-      (AF.Empty _)       -> error "tree must have at least one block"
-      (_ AF.:> tipBlock) -> tipBlock
-
     -- 1. The adversary advertises blocks up to the intersection.
     -- 2. The honest node advertises all its chain, which is
     --    long enough to be blocked by the LoE.
@@ -505,9 +496,7 @@ prop_densityDisconnectTriggersChainSel =
           intersect = case btbPrefix branch of
             (AF.Empty _)       -> Origin
             (_ AF.:> tipBlock) -> At tipBlock
-          advTip = case btbFull branch of
-            (AF.Empty _) -> error "alternate branch must have at least one block"
-            (_ AF.:> tipBlock) -> tipBlock
+          advTip = getOnlyBranchTip tree
        in peers'
             -- Eagerly serve the honest tree, but after the adversary has
             -- advertised its chain up to the intersection.

ouroboros-consensus-diffusion/test/consensus-test/Test/Consensus/Genesis/Tests/LoE.hs

@@ -27,6 +27,7 @@ import           Test.Consensus.PointSchedule.SinglePeer (scheduleBlockPoint,
 import           Test.Tasty
 import           Test.Tasty.QuickCheck
 import           Test.Util.Orphans.IOLike ()
+import           Test.Util.PartialAccessors
 import           Test.Util.TestEnv (adjustQuickCheckTests)
 
 tests :: TestTree
@@ -77,23 +78,14 @@ prop_adversaryHitsTimeouts timeoutsEnabled =
            in selectedCorrect && exceptionsCorrect
       )
   where
-    getOnlyBranch :: BlockTree blk -> BlockTreeBranch blk
-    getOnlyBranch BlockTree {btBranches} = case btBranches of
-      [branch] -> branch
-      _        -> error "tree must have exactly one alternate branch"
-
     delaySchedule :: HasHeader blk => BlockTree blk -> Peers (PeerSchedule blk)
     delaySchedule tree =
-      let trunkTip = case btTrunk tree of
-            (AF.Empty _)       -> error "tree must have at least one block"
-            (_ AF.:> tipBlock) -> tipBlock
+      let trunkTip = getTrunkTip tree
           branch = getOnlyBranch tree
           intersectM = case btbPrefix branch of
             (AF.Empty _)       -> Nothing
             (_ AF.:> tipBlock) -> Just tipBlock
-          branchTip = case btbFull branch of
-            (AF.Empty _) -> error "alternate branch must have at least one block"
-            (_ AF.:> tipBlock) -> tipBlock
+          branchTip = getOnlyBranchTip tree
        in peers'
             -- Eagerly serve the honest tree, but after the adversary has
             -- advertised its chain.

ouroboros-consensus-diffusion/test/consensus-test/Test/Consensus/Genesis/Tests/LoP.hs

@@ -31,6 +31,7 @@ import           Test.Consensus.PointSchedule.SinglePeer (scheduleBlockPoint,
 import           Test.Tasty
 import           Test.Tasty.QuickCheck
 import           Test.Util.Orphans.IOLike ()
+import           Test.Util.PartialAccessors
 import           Test.Util.TestEnv (adjustQuickCheckTests)
 
 tests :: TestTree
@@ -213,23 +214,14 @@ prop_delayAttack lopEnabled =
            in selectedCorrect && exceptionsCorrect
       )
   where
-    getOnlyBranch :: BlockTree blk -> BlockTreeBranch blk
-    getOnlyBranch BlockTree {btBranches} = case btBranches of
-      [branch] -> branch
-      _        -> error "tree must have exactly one alternate branch"
-
     delaySchedule :: (HasHeader blk) => BlockTree blk -> Peers (PeerSchedule blk)
     delaySchedule tree =
-      let trunkTip = case btTrunk tree of
-            (AF.Empty _)       -> error "tree must have at least one block"
-            (_ AF.:> tipBlock) -> tipBlock
+      let trunkTip = getTrunkTip tree
           branch = getOnlyBranch tree
           intersectM = case btbPrefix branch of
             (AF.Empty _)       -> Nothing
             (_ AF.:> tipBlock) -> Just tipBlock
-          branchTip = case btbFull branch of
-            (AF.Empty _) -> error "alternate branch must have at least one block"
-            (_ AF.:> tipBlock) -> tipBlock
+          branchTip = getOnlyBranchTip tree
        in peers'
             -- Eagerly serve the honest tree, but after the adversary has
             -- advertised its chain.

ouroboros-consensus-diffusion/test/consensus-test/Test/Util/PartialAccessors.hs

@@ -0,0 +1,32 @@
+{-# LANGUAGE NamedFieldPuns #-}
+
+-- | Helpers to access particular parts of trees and schedules
+-- Those functions are partial, and are designed to only be used in tests.
+-- We know they won't fail there, because we generated the structures
+-- with the correct properties.
+module Test.Util.PartialAccessors (
+    getOnlyBranch
+  , getOnlyBranchTip
+  , getTrunkTip
+  ) where
+
+import qualified Ouroboros.Network.AnchoredFragment as AF
+import           Ouroboros.Network.Block (HasHeader)
+import           Test.Consensus.BlockTree
+
+getOnlyBranch :: BlockTree blk -> BlockTreeBranch blk
+getOnlyBranch BlockTree {btBranches} = case btBranches of
+  [branch] -> branch
+  _        -> error "tree must have exactly one alternate branch"
+
+getTrunkTip :: HasHeader blk => BlockTree blk -> blk
+getTrunkTip tree = case btTrunk tree of
+  (AF.Empty _)       -> error "tree must have at least one block"
+  (_ AF.:> tipBlock) -> tipBlock
+
+getOnlyBranchTip :: HasHeader blk => BlockTree blk -> blk
+getOnlyBranchTip BlockTree {btBranches} = case btBranches of
+  [branch] -> case btbFull branch of
+    (AF.Empty _)       -> error "alternate branch must have at least one block"
+    (_ AF.:> tipBlock) -> tipBlock
+  _ -> error "tree must have exactly one alternate branch"