Fixed and extended groupConns

IntersectMBO · Apr 20, 2022 · c1d0c10 · c1d0c10
1 parent f9215ea
commit c1d0c10
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 14 deletions.
diff --git a/ouroboros-network/test/Test/Ouroboros/Network/Testnet.hs b/ouroboros-network/test/Test/Ouroboros/Network/Testnet.hs
@@ -45,6 +45,7 @@ import           Ouroboros.Network.Diffusion.P2P
                       (TracersExtra(..), RemoteTransitionTrace)
 import           Ouroboros.Network.ConnectionHandler (ConnectionHandlerTrace)
 import           Ouroboros.Network.ConnectionManager.Types
+import           Ouroboros.Network.InboundGovernor (RemoteSt)
 import qualified Ouroboros.Network.Testing.Data.Signal as Signal
 import qualified Ouroboros.Network.PeerSelection.Governor as Governor
 import qualified Ouroboros.Network.PeerSelection.LocalRootPeers as LocalRootPeers
@@ -884,6 +885,10 @@ prop_diffusion_cm_valid_transitions defaultBearerInfo diffScript =
       <$> events
 
   where
+    isFinalTransition :: Transition' AbstractState -> Bool
+    isFinalTransition (Transition _ UnknownConnectionSt) = True
+    isFinalTransition _                                  = False
+
     verify_cm_valid_transitions :: Trace () DiffusionTestTrace -> Property
     verify_cm_valid_transitions events =
       let abstractTransitionEvents :: Trace () (AbstractTransitionTrace NtNAddr)
@@ -932,7 +937,7 @@ prop_diffusion_cm_valid_transitions defaultBearerInfo diffScript =
               }
            )
         . fmap (map ttTransition)
-        . groupConns id
+        . groupConns id isFinalTransition
         $ abstractTransitionEvents
 
 
@@ -986,6 +991,10 @@ prop_diffusion_cm_valid_transition_order defaultBearerInfo diffScript =
         )
       <$> events
   where
+    isFinalTransition :: Transition' AbstractState -> Bool
+    isFinalTransition (Transition _ UnknownConnectionSt) = True
+    isFinalTransition _                                  = False
+
     verify_cm_valid_transition_order :: Trace () DiffusionTestTrace -> Property
     verify_cm_valid_transition_order events =
       let abstractTransitionEvents :: Trace () (AbstractTransitionTrace NtNAddr)
@@ -997,7 +1006,7 @@ prop_diffusion_cm_valid_transition_order defaultBearerInfo diffScript =
             (const mempty)
             verifyAbstractTransitionOrder
          . fmap (map ttTransition)
-         . groupConns id
+         . groupConns id isFinalTransition
          $ abstractTransitionEvents
 
 -- | A variant of ouroboros-network-framework
@@ -1123,6 +1132,10 @@ prop_diffusion_ig_valid_transition_order defaultBearerInfo diffScript =
       <$> events
 
   where
+    isFinalTransition :: Transition' (Maybe RemoteSt) -> Bool
+    isFinalTransition (Transition _ Nothing) = True
+    isFinalTransition _                      = False
+
     verify_ig_valid_transition_order :: Trace () DiffusionTestTrace -> Property
     verify_ig_valid_transition_order events =
 
@@ -1135,7 +1148,7 @@ prop_diffusion_ig_valid_transition_order defaultBearerInfo diffScript =
            (const mempty)
            verifyRemoteTransitionOrder
         . fmap (map ttTransition)
-        . groupConns id
+        . groupConns id isFinalTransition
         $ remoteTransitionTraceEvents
 
 -- | A variant of ouroboros-network-framework
@@ -1192,10 +1205,14 @@ prop_diffusion_timeouts_enforced defaultBearerInfo diffScript =
       <$> events
 
   where
+    isFinalTransition :: Transition' AbstractState -> Bool
+    isFinalTransition (Transition _ UnknownConnectionSt) = True
+    isFinalTransition _                                  = False
+
     verify_timeouts :: Trace () (Time, DiffusionTestTrace) -> Property
     verify_timeouts events =
       let transitionSignal :: Trace () [(Time, AbstractTransitionTrace NtNAddr)]
-          transitionSignal = groupConns snd
+          transitionSignal = groupConns snd isFinalTransition
                            . selectDiffusionConnectionManagerTransitionEventsTime
                            $ events
 

diff --git a/ouroboros-network/test/Test/Ouroboros/Network/Testnet/Simulation/ConnectionManager.hs b/ouroboros-network/test/Test/Ouroboros/Network/Testnet/Simulation/ConnectionManager.hs
@@ -73,24 +73,32 @@ splitConns getTransition =
 --
 groupConns :: Ord addr
            => (a -> TransitionTrace' addr st)
+           -> (Transition' st -> Bool)
            -> Trace r a
            -> Trace r [a]
-groupConns getTransition =
+groupConns getTransition isFinalTransition =
     fmap fromJust
   . Trace.filter isJust
   -- there might be some connections in the state, push them onto the 'Trace'
-  . (\(s, o) -> foldr (\a as -> Trace.Cons (Just a) as) o (Map.elems s))
+  . (\(s, o) -> foldr (\a as -> Trace.Cons (Just (reverse a)) as) o (Map.elems s))
   . bimapAccumL
       ( \ s a -> (s, a))
       ( \ s a ->
-          let TransitionTrace { ttPeerAddr } = getTransition a
-           in case ttPeerAddr `Map.lookup` s of
-                Nothing  -> ( Map.insert ttPeerAddr [a] s
-                            , Nothing
-                            )
-                Just trs -> ( Map.delete ttPeerAddr s
-                            , Just (reverse $ a : trs)
-                            )
+          let TransitionTrace { ttPeerAddr, ttTransition } = getTransition a
+           in if isFinalTransition ttTransition
+                 then case ttPeerAddr `Map.lookup` s of
+                        Nothing  -> ( Map.insert ttPeerAddr [a] s
+                                    , Nothing
+                                    )
+                        Just trs -> ( Map.delete ttPeerAddr s
+                                    , Just (reverse $ a : trs)
+                                    )
+                 else ( Map.alter (\case
+                                     Nothing -> Just [a]
+                                     Just as -> Just (a : as)
+                                  )
+                                  ttPeerAddr s
+                      , Nothing)
       )
       Map.empty