1 files changed, 29 insertions, 19 deletions
diff --git a/src/Test.hs b/src/Test.hs
index 3458c04..c2a35e8 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -6,6 +6,7 @@ module Test (
 
 import Data.Scientific
 import Data.Text (Text)
+import Data.Typeable
 
 import Network
 import Process
@@ -13,26 +14,35 @@ import Script.Expr
 
 data Test = Test
     { testName :: Text
-    , testSteps :: Expr TestBlock
+    , testSteps :: Expr (TestBlock ())
     }
 
-newtype TestBlock = TestBlock [ TestStep ]
-    deriving (Semigroup, Monoid)
-
-data TestStep
-    = Subnet (TypedVarName Network) Network (Network -> TestBlock)
-    | DeclNode (TypedVarName Node) Network (Node -> TestBlock)
-    | Spawn (TypedVarName Process) (Either Network Node) (Process -> TestBlock)
-    | Send Process Text
-    | Expect SourceLine Process (Traced Regex) [ TypedVarName Text ] ([ Text ] -> TestBlock)
-    | Flush Process (Maybe Regex)
-    | Guard SourceLine EvalTrace Bool
-    | DisconnectNode Node TestBlock
-    | DisconnectNodes Network TestBlock
-    | DisconnectUpstream Network TestBlock
-    | PacketLoss Scientific Node TestBlock
-    | Wait
-
-instance ExprType TestBlock where
+data TestBlock a where
+    EmptyTestBlock :: TestBlock ()
+    TestBlockStep :: TestBlock () -> TestStep a -> TestBlock a
+
+instance Semigroup (TestBlock ()) where
+    EmptyTestBlock <> block = block
+    block <> EmptyTestBlock = block
+    block <> TestBlockStep block' step = TestBlockStep (block <> block') step
+
+instance Monoid (TestBlock ()) where
+    mempty = EmptyTestBlock
+
+data TestStep a where
+    Subnet :: TypedVarName Network -> Network -> (Network -> TestBlock a) -> TestStep a
+    DeclNode :: TypedVarName Node -> Network -> (Node -> TestBlock a) -> TestStep a
+    Spawn :: TypedVarName Process -> Either Network Node -> (Process -> TestBlock a) -> TestStep a
+    Send :: Process -> Text -> TestStep ()
+    Expect :: SourceLine -> Process -> Traced Regex -> [ TypedVarName Text ] -> ([ Text ] -> TestBlock a) -> TestStep a
+    Flush :: Process -> Maybe Regex -> TestStep ()
+    Guard :: SourceLine -> EvalTrace -> Bool -> TestStep ()
+    DisconnectNode :: Node -> TestBlock a -> TestStep a
+    DisconnectNodes :: Network -> TestBlock a -> TestStep a
+    DisconnectUpstream :: Network -> TestBlock a -> TestStep a
+    PacketLoss :: Scientific -> Node -> TestBlock a -> TestStep a
+    Wait :: TestStep ()
+
+instance Typeable a => ExprType (TestBlock a) where
     textExprType _ = "test block"
     textExprValue _ = "<test block>"