Assign regex captures to variables

3 files changed, 40 insertions, 11 deletions

diff --git a/src/Main.hs b/src/Main.hs
index 3fa3468..a1b6625 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -263,13 +263,13 @@ getProcess net pname = liftIO $ do
     Just p <- find ((pname==).procName) <$> readMVar (netProcesses net)
     return p
 
-tryMatch :: Regex -> [Text] -> Maybe (Text, [Text])
-tryMatch re (x:xs) | Right (Just _) <- regexec re x = Just (x, xs)
+tryMatch :: Regex -> [Text] -> Maybe ((Text, [Text]), [Text])
+tryMatch re (x:xs) | Right (Just (_, _, _, capture)) <- regexec re x = Just ((x, capture), xs)
                    | otherwise = fmap (x:) <$> tryMatch re xs
 tryMatch _ [] = Nothing
 
-expect :: Process -> Regex -> Text -> TestRun ()
-expect p re pat = do
+expect :: Process -> Regex -> Text -> [VarName] -> TestRun ()
+expect p re pat vars = do
     timeout <- asks $ optTimeout . teOptions
     delay <- liftIO $ registerDelay $ ceiling $ 1000000 * timeout
     mbmatch <- atomicallyTest $ (Nothing <$ (check =<< readTVar delay)) <|> do
@@ -280,7 +280,18 @@ expect p re pat = do
                  writeTVar (procOutput p) out'
                  return $ Just m
     case mbmatch of
-         Just line -> do
+         Just (line, capture) -> do
+             when (length vars /= length capture) $ do
+                 outLine OutputMatchFail (Just $ procName p) $ T.pack "mismatched number of capture variables /" `T.append` pat `T.append` T.pack "/"
+                 throwError ()
+
+             forM_ vars $ \name -> do
+                 cur <- gets (lookup name . tsVars)
+                 when (isJust cur) $ do
+                     outLine OutputMatchFail (Just $ procName p) $ T.pack "variable already exists: '" `T.append` textVarName name `T.append` T.pack "'"
+                     throwError ()
+
+             modify $ \s -> s { tsVars = zip vars capture ++ tsVars s }
              outLine OutputMatch (Just $ procName p) line
          Nothing -> do
              outLine OutputMatchFail (Just $ procName p) $ T.pack "expect failed /" `T.append` pat `T.append` T.pack "/"
@@ -325,11 +336,11 @@ runTest out opts test = do
                 line <- evalStringExpr expr
                 send p line
 
-            Expect pname expr@(RegexExpr ps) -> do
+            Expect pname expr@(RegexExpr ps) captures -> do
                 p <- getProcess net pname
                 regex <- evalRegexExpr expr
                 pat <- evalStringExpr (StringExpr $ map (left T.pack) ps)
-                expect p regex pat
+                expect p regex pat captures
 
             Wait -> do
                 outPrompt $ T.pack "Waiting..."
diff --git a/src/Parser.hs b/src/Parser.hs
index ddbdfd6..bf7f75b 100644
--- a/src/Parser.hs
+++ b/src/Parser.hs
@@ -12,6 +12,7 @@ import Control.Monad.State
 import Data.Char
 import Data.Set (Set)
 import qualified Data.Set as S
+import Data.Text (Text)
 import Data.Text qualified as T
 import qualified Data.Text.Lazy as TL
 import qualified Data.Text.Lazy.IO as TL
@@ -68,6 +69,13 @@ block merge header item = L.indentBlock scn $ do
         , L.IndentNone <$> merge h []
         ]
 
+listOf :: TestParser a -> TestParser [a]
+listOf item = do
+    sc
+    x <- item
+    sc
+    (x:) <$> choice [ char ',' >> listOf item, return [] ]
+
 nodeName :: TestParser NodeName
 nodeName = label "network node name" $ lexeme $ do
     c <- lowerChar
@@ -80,11 +88,15 @@ procName = label "process name" $ lexeme $ do
     cs <- takeWhileP Nothing (\x -> isAlphaNum x || x == '_' || x == '-')
     return $ ProcName $ TL.toStrict (c `TL.cons` cs)
 
+identifier :: TestParser Text
+identifier = do
+    TL.toStrict <$> takeWhile1P Nothing (\x -> isAlphaNum x || x == '_')
+
 varExpansion :: TestParser VarName
 varExpansion = do
     void $ char '$'
     choice
-        [ VarName . (:[]) . TL.toStrict <$> takeWhile1P Nothing (\x -> isAlphaNum x || x == '_')
+        [ VarName . (:[]) <$> identifier
         ,do void $ char '{'
             name <- takeWhile1P Nothing (/='}')
             void $ char '}'
@@ -206,6 +218,7 @@ testSend = command "send"
 data ExpectBuilder = ExpectBuilder
     { _expectBuilderProc :: Maybe ProcName
     , _expectBuilderRegex :: Maybe RegexExpr
+    , _expectBuilderCaptures :: Maybe [VarName]
     }
     deriving (Generic)
 
@@ -215,9 +228,11 @@ testExpect :: TestParser [TestStep]
 testExpect = command "expect"
     [ Param "from" expectBuilderProc procName
     , Param "" expectBuilderRegex regex
+    , Param "capture" expectBuilderCaptures (listOf $ VarName . (:[]) <$> identifier)
     ] $ \b -> Expect
         <$> (maybe (fail "missing 'from' <proc>") return $ b ^. expectBuilderProc)
         <*> (maybe (fail "missing regex to match") return $ b ^. expectBuilderRegex)
+        <*> (maybe (return []) return $ b ^. expectBuilderCaptures)
 
 
 testWait :: TestParser [TestStep]
diff --git a/src/Test.hs b/src/Test.hs
index d652f9b..c58a2a3 100644
--- a/src/Test.hs
+++ b/src/Test.hs
@@ -6,7 +6,7 @@ module Test (
     NodeName(..), textNodeName, unpackNodeName,
 
     MonadEval(..),
-    VarName(..), unpackVarName,
+    VarName(..), textVarName, unpackVarName,
     StringExpr(..), evalStringExpr,
     RegexExpr(..), evalRegexExpr,
 ) where
@@ -29,7 +29,7 @@ data Test = Test
 
 data TestStep = Spawn ProcName NodeName
               | Send ProcName StringExpr
-              | Expect ProcName RegexExpr
+              | Expect ProcName RegexExpr [VarName]
               | Wait
 
 newtype NodeName = NodeName Text
@@ -49,8 +49,11 @@ class Monad m => MonadEval m where
 data VarName = VarName [Text]
     deriving (Eq, Ord)
 
+textVarName :: VarName -> Text
+textVarName (VarName name) = T.concat $ intersperse (T.singleton '.') name
+
 unpackVarName :: VarName -> String
-unpackVarName (VarName name) = concat $ intersperse "." $ map T.unpack name
+unpackVarName = T.unpack . textVarName
 
 data StringExpr = StringExpr [Either Text VarName]