Explicit Scope constructor in TestStep data type

4 files changed, 44 insertions, 35 deletions

diff --git a/src/Parser.hs b/src/Parser.hs
index 0716457..9f1a0e3 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -43,7 +43,7 @@ parseTestDefinition = label "test definition" $ toplevel ToplevelTest $ do
         modify $ \s -> s
             { testContext = SomeExpr $ varExpr SourceLineBuiltin rootNetworkVar
             }
-        block (\name steps -> return $ Test name $ mconcat steps) header testStep
+        block (\name steps -> return $ Test name $ Scope <$> mconcat steps) header testStep
   where
     header = do
         wsymbol "test"
@@ -64,7 +64,7 @@ parseDefinition href = label "symbol definition" $ do
                 osymbol ":"
                 scn
                 ref <- L.indentGuard scn GT href
-                SomeExpr <$> blockOf ref testStep
+                SomeExpr <$> testBlock ref
             , do
                 osymbol "="
                 someExpr <* eol
diff --git a/src/Parser/Statement.hs b/src/Parser/Statement.hs
index 27e7b92..474fa03 100644
--- a/src/Parser/Statement.hs
+++ b/src/Parser/Statement.hs
@@ -1,5 +1,6 @@
 module Parser.Statement (
     testStep,
+    testBlock,
 ) where
 
 import Control.Monad
@@ -43,7 +44,7 @@ letStatement = do
         addVarName off tname
         void $ eol
         body <- testBlock indent
-        return $ Let line tname e body
+        return $ Let line tname e (TestBlockStep EmptyTestBlock . Scope <$> body)
 
 forStatement :: TestParser (Expr (TestBlock ()))
 forStatement = do
@@ -68,7 +69,7 @@ forStatement = do
         body <- testBlock indent
         return $ (\xs f -> mconcat $ map f xs)
             <$> (unpack <$> e)
-            <*> LambdaAbstraction tname body
+            <*> LambdaAbstraction tname (TestBlockStep EmptyTestBlock . Scope <$> body)
 
 shellStatement :: TestParser (Expr (TestBlock ()))
 shellStatement = do
@@ -108,7 +109,7 @@ shellStatement = do
             void eol
             void $ L.indentGuard scn GT ref
             script <- shellScript
-            cont <- testBlock ref
+            cont <- fmap Scope <$> testBlock ref
             let expr | Just pname <- mbpname = LambdaAbstraction pname cont
                      | otherwise = const <$> cont
             return $ TestBlockStep EmptyTestBlock <$>
@@ -290,14 +291,14 @@ instance ExprType a => ParamType (InnerBlock a) where
         combine f (x : xs) = f x xs
         combine _ [] = error "inner block parameter count mismatch"
 
-innerBlock :: CommandDef (TestBlock ())
+innerBlock :: CommandDef (TestStep ())
 innerBlock = ($ ([] :: [ Void ])) <$> innerBlockFun
 
-innerBlockFun :: ExprType a => CommandDef (a -> TestBlock ())
+innerBlockFun :: ExprType a => CommandDef (a -> TestStep ())
 innerBlockFun = (\f x -> f [ x ]) <$> innerBlockFunList
 
-innerBlockFunList :: ExprType a => CommandDef ([ a ] -> TestBlock ())
-innerBlockFunList = fromInnerBlock <$> param ""
+innerBlockFunList :: ExprType a => CommandDef ([ a ] -> TestStep ())
+innerBlockFunList = (\ib -> Scope . fromInnerBlock ib) <$> param ""
 
 newtype ExprParam a = ExprParam { fromExprParam :: a }
     deriving (Functor, Foldable, Traversable)
@@ -380,7 +381,8 @@ testLocal = do
     void $ eol
 
     indent <- L.indentGuard scn GT ref
-    localState $ testBlock indent
+    localState $ do
+        fmap (TestBlockStep EmptyTestBlock . Scope) <$> testBlock indent
 
 testWith :: TestParser (Expr (TestBlock ()))
 testWith = do
@@ -406,7 +408,7 @@ testWith = do
     indent <- L.indentGuard scn GT ref
     localState $ do
         modify $ \s -> s { testContext = ctx }
-        testBlock indent
+        fmap (TestBlockStep EmptyTestBlock . Scope) <$> testBlock indent
 
 testSubnet :: TestParser (Expr (TestBlock ()))
 testSubnet = command "subnet" $ Subnet
diff --git a/src/Run.hs b/src/Run.hs
index 200ae8e..32b04c6 100644
--- a/src/Run.hs
+++ b/src/Run.hs
@@ -96,7 +96,7 @@ runTest out opts gdefs test = do
     resetOutputTime out
     res <- runExceptT $ flip runReaderT (tenv, tstate) $ fromTestRun $ do
         withInternet $ \_ -> do
-            evalBlock =<< eval (testSteps test)
+            runStep =<< eval (testSteps test)
             when (optWait opts) $ do
                 void $ outPromptGetLine $ "Test '" <> testName test <> "' completed, waiting..."
 
@@ -137,14 +137,20 @@ evalGlobalDefs :: [ (( ModuleName, VarName ), SomeExpr ) ] -> GlobalDefs
 evalGlobalDefs exprs = fix $ \gdefs ->
     builtins `M.union` M.fromList (map (fmap (evalSomeWith gdefs)) exprs)
 
-evalBlock :: TestBlock () -> TestRun ()
-evalBlock EmptyTestBlock = return ()
-evalBlock (TestBlockStep prev step) = evalBlock prev >> case step of
+runBlock :: TestBlock () -> TestRun ()
+runBlock EmptyTestBlock = return ()
+runBlock (TestBlockStep prev step) = runBlock prev >> runStep step
+
+runStep :: TestStep () -> TestRun ()
+runStep = \case
+    Scope block -> do
+        runBlock block
+
     Subnet name parent inner -> do
-        withSubnet parent (Just name) $ evalBlock . inner
+        withSubnet parent (Just name) $ runStep . inner
 
     DeclNode name net inner -> do
-        withNode net (Left name) $ evalBlock . inner
+        withNode net (Left name) $ runStep . inner
 
     Spawn tvname@(TypedVarName (VarName tname)) target args inner -> do
         case target of
@@ -157,20 +163,20 @@ evalBlock (TestBlockStep prev step) = evalBlock prev >> case step of
                 tool = fromMaybe (optDefaultTool opts) (lookup pname $ optProcTools opts)
                 cmd = unwords $ tool : map (T.unpack . escape) args
                 escape = ("'" <>) . (<> "'") . T.replace "'" "'\\''"
-            withProcess (Right node) pname Nothing cmd $ evalBlock . inner
+            withProcess (Right node) pname Nothing cmd $ runStep . inner
 
     SpawnShell mbname node script inner -> do
         let tname | Just (TypedVarName (VarName name)) <- mbname = name
                   | otherwise = "shell"
         let pname = ProcName tname
-        withShellProcess node pname script $ evalBlock . inner
+        withShellProcess node pname script $ runStep . inner
 
     Send p line -> do
         outProc OutputChildStdin p line
         send p line
 
     Expect line p expr captures inner -> do
-        expect line p expr captures $ evalBlock . inner
+        expect line p expr captures $ runStep . inner
 
     Flush p regex -> do
         flush p regex
@@ -179,18 +185,18 @@ evalBlock (TestBlockStep prev step) = evalBlock prev >> case step of
         testStepGuard line vars expr
 
     DisconnectNode node inner -> do
-        withDisconnectedUp (nodeUpstream node) $ evalBlock inner
+        withDisconnectedUp (nodeUpstream node) $ runStep inner
 
     DisconnectNodes net inner -> do
-        withDisconnectedBridge (netBridge net) $ evalBlock inner
+        withDisconnectedBridge (netBridge net) $ runStep inner
 
     DisconnectUpstream net inner -> do
         case netUpstream net of
-            Just link -> withDisconnectedUp link $ evalBlock inner
-            Nothing -> evalBlock inner
+            Just link -> withDisconnectedUp link $ runStep inner
+            Nothing -> runStep inner
 
     PacketLoss loss node inner -> do
-        withNodePacketLoss node loss $ evalBlock inner
+        withNodePacketLoss node loss $ runStep inner
 
     Wait -> do
         void $ outPromptGetLine "Waiting..."
diff --git a/src/Test.hs b/src/Test.hs
index ff51ebe..a9a2cdb 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -15,7 +15,7 @@ import Script.Shell
 
 data Test = Test
     { testName :: Text
-    , testSteps :: Expr (TestBlock ())
+    , testSteps :: Expr (TestStep ())
     }
 
 data TestBlock a where
@@ -31,18 +31,19 @@ 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 -> [ Text ] -> (Process -> TestBlock a) -> TestStep a
-    SpawnShell :: Maybe (TypedVarName Process) -> Node -> ShellScript -> (Process -> TestBlock a) -> TestStep a
+    Scope :: TestBlock a -> TestStep a
+    Subnet :: TypedVarName Network -> Network -> (Network -> TestStep a) -> TestStep a
+    DeclNode :: TypedVarName Node -> Network -> (Node -> TestStep a) -> TestStep a
+    Spawn :: TypedVarName Process -> Either Network Node -> [ Text ] -> (Process -> TestStep a) -> TestStep a
+    SpawnShell :: Maybe (TypedVarName Process) -> Node -> ShellScript -> (Process -> TestStep a) -> TestStep a
     Send :: Process -> Text -> TestStep ()
-    Expect :: SourceLine -> Process -> Traced Regex -> [ TypedVarName Text ] -> ([ Text ] -> TestBlock a) -> TestStep a
+    Expect :: SourceLine -> Process -> Traced Regex -> [ TypedVarName Text ] -> ([ Text ] -> TestStep 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
+    DisconnectNode :: Node -> TestStep a -> TestStep a
+    DisconnectNodes :: Network -> TestStep a -> TestStep a
+    DisconnectUpstream :: Network -> TestStep a -> TestStep a
+    PacketLoss :: Scientific -> Node -> TestStep a -> TestStep a
     Wait :: TestStep ()
 
 instance Typeable a => ExprType (TestBlock a) where