Test data type

2 files changed, 133 insertions, 61 deletions

diff --git a/src/Main.hs b/src/Main.hs
index fbd9bea..ca501bf 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -4,6 +4,12 @@ import Control.Concurrent
 import Control.Concurrent.STM
 import Control.Monad
 
+import Data.List
+import Data.Maybe
+import Data.Text (Text)
+import qualified Data.Text as T
+import qualified Data.Text.IO as T
+
 import Text.Regex.TDFA
 import Text.Regex.TDFA.String
 
@@ -15,20 +21,23 @@ import System.IO
 import System.IO.Error
 import System.Process
 
+import Test
+
 data Network = Network
-    { netNodes :: MVar [(Int, Node)]
+    { netNodes :: MVar [Node]
     , netDir :: FilePath
     }
 
 data Node = Node
-    { nodeNetwork :: Network
+    { nodeName :: NodeName
+    , nodeNetwork :: Network
     , nodeProcesses :: MVar [Process]
-    , nodeName :: String
     , nodeDir :: FilePath
     }
 
 data Process = Process
-    { procHandle :: ProcessHandle
+    { procName :: ProcName
+    , procHandle :: ProcessHandle
     , procNode :: Node
     , procStdin :: Handle
     , procOutput :: TVar [String]
@@ -52,14 +61,14 @@ initNetwork = do
 exitNetwork :: Network -> IO ()
 exitNetwork net = do
     nodes <- readMVar (netNodes net)
-    ok <- fmap and $ forM nodes $ \(_, node) -> do
+    ok <- fmap and $ forM nodes $ \node -> do
         processes <- readMVar (nodeProcesses node)
         fmap and $ forM processes $ \p -> do
             hClose (procStdin p)
             waitForProcess (procHandle p) >>= \case
                 ExitSuccess -> return True
                 ExitFailure code -> do
-                    putStrLn $ "\ESC[31m" ++ nodeName node ++ "!!> exit code: " ++ show code ++ "\ESC[0m"
+                    putStrLn $ "\ESC[31m" ++ unpackNodeName (nodeName node) ++ "!!> exit code: " ++ show code ++ "\ESC[0m"
                     return False
 
     if ok
@@ -67,16 +76,16 @@ exitNetwork net = do
                exitSuccess
        else exitFailure
 
-getNode :: Network -> Int -> IO Node
-getNode net idx = (lookup idx <$> readMVar (netNodes net)) >>= \case
+getNode :: Network -> NodeName -> IO Node
+getNode net nname@(NodeName tnname) = (find ((nname==).nodeName) <$> readMVar (netNodes net)) >>= \case
     Just node -> return node
     _ -> do
         processes <- newMVar []
-        let name = "node" ++ show idx
-            dir = netDir net </> ("erebos" ++ show idx)
-            node = Node { nodeNetwork = net
+        let name = T.unpack tnname
+            dir = netDir net </> ("erebos_" ++ name)
+            node = Node { nodeName = nname
+                        , nodeNetwork = net
                         , nodeProcesses = processes
-                        , nodeName = name
                         , nodeDir = dir
                         }
 
@@ -84,21 +93,22 @@ getNode net idx = (lookup idx <$> readMVar (netNodes net)) >>= \case
         when exists $ ioError $ userError $ dir ++ " exists"
         createDirectoryIfMissing True dir
 
-        callCommand $ "ip netns add \""++ name ++ "\""
-        callCommand $ "ip link add \"veth" ++ show idx ++ ".0\" type veth peer name \"veth" ++ show idx ++ ".1\" netns \"" ++ name ++ "\""
-        callCommand $ "ip link set dev \"veth" ++ show idx ++ ".0\" master br0 up"
-        callOn node $ "ip addr add 192.168.0." ++ show (10 + idx) ++ "/24 broadcast 192.168.0.255 dev \"veth" ++ show idx ++ ".1\""
-        callOn node $ "ip link set dev \"veth" ++ show idx ++ ".1\" up"
-        callOn node $ "ip link set dev lo up"
-        modifyMVar_ (netNodes net) $ return . ((idx, node):)
+        modifyMVar_ (netNodes net) $ \nodes -> do
+            callCommand $ "ip netns add \""++ name ++ "\""
+            callCommand $ "ip link add \"veth_" ++ name ++ ".0\" type veth peer name \"veth_" ++ name ++ ".1\" netns \"" ++ name ++ "\""
+            callCommand $ "ip link set dev \"veth_" ++ name ++ ".0\" master br0 up"
+            callOn node $ "ip addr add 192.168.0." ++ show (11 + length nodes) ++ "/24 broadcast 192.168.0.255 dev \"veth_" ++ name ++ ".1\""
+            callOn node $ "ip link set dev \"veth_" ++ name++ ".1\" up"
+            callOn node $ "ip link set dev lo up"
+            return $ node : nodes
         return node
 
 callOn :: Node -> String -> IO ()
-callOn node cmd = callCommand $ "ip netns exec \"" ++ nodeName node ++ "\" " ++ cmd
+callOn node cmd = callCommand $ "ip netns exec \"" ++ unpackNodeName (nodeName node) ++ "\" " ++ cmd
 
-spawnOn :: Node -> String -> IO Process
-spawnOn node cmd = do
-    (Just hin, Just hout, Just herr, handle) <- createProcess (shell $ "ip netns exec \"" ++ nodeName node ++ "\" " ++ cmd)
+spawnOn :: Node -> ProcName -> String -> IO Process
+spawnOn node pname cmd = do
+    (Just hin, Just hout, Just herr, handle) <- createProcess (shell $ "ip netns exec \"" ++ unpackNodeName (nodeName node) ++ "\" " ++ cmd)
         { std_in = CreatePipe, std_out = CreatePipe, std_err = CreatePipe
         , env = Just [("EREBOS_DIR", nodeDir node)]
         }
@@ -109,19 +119,20 @@ spawnOn node cmd = do
             tryIOError (hGetLine h) >>= \case
                 Left err
                     | isEOFError err -> return ()
-                    | otherwise -> putStrLn $ "\ESC[31m" ++ nodeName node ++ "!!> IO error: " ++ show err ++ "\ESC[0m"
+                    | otherwise -> putStrLn $ "\ESC[31m" ++ unpackNodeName (nodeName node) ++ "!!> IO error: " ++ show err ++ "\ESC[0m"
                 Right line -> do
                     act line
                     readingLoop h act
 
     void $ forkIO $ readingLoop hout $ \line -> do
-        putStrLn $ nodeName node ++ "> " ++ line
+        putStrLn $ unpackNodeName (nodeName node) ++ "> " ++ line
         atomically $ modifyTVar out (++[line])
     void $ forkIO $ readingLoop herr $ \line -> do
-        putStrLn $ "\ESC[31m" ++ nodeName node ++ "!> " ++ line ++ "\ESC[0m"
+        putStrLn $ "\ESC[31m" ++ unpackNodeName (nodeName node) ++ "!> " ++ line ++ "\ESC[0m"
 
     let process = Process
-            { procHandle = handle
+            { procName = pname
+            , procHandle = handle
             , procNode = node
             , procStdin = hin
             , procOutput = out
@@ -130,49 +141,77 @@ spawnOn node cmd = do
     modifyMVar_ (nodeProcesses node) $ return . (process:)
     return process
 
+getProcess :: Network -> ProcName -> IO Process
+getProcess net pname = do
+    nodes <- readMVar (netNodes net)
+    (p:_) <- fmap catMaybes $ forM nodes $ \node -> do
+        processes <- readMVar (nodeProcesses node)
+        return $ find ((pname==).procName) processes
+    return p
+
 tryMatch :: Regex -> [String] -> Maybe (String, [String])
 tryMatch re (x:xs) | Right (Just _) <- regexec re x = Just (x, xs)
                    | otherwise = fmap (x:) <$> tryMatch re xs
 tryMatch _ [] = Nothing
 
-expect :: Process -> String -> IO ()
-expect p pat = case compile defaultCompOpt defaultExecOpt ("^" ++ pat ++ "$") of
-                    Right re -> do
-                        mbmatch <- atomically $ do
-                            out <- readTVar (procOutput p)
-                            case tryMatch re out of
-                                 Nothing -> retry
-                                 Just (m, out') -> do
-                                     writeTVar (procOutput p) out'
-                                     return $ Just m
-                        case mbmatch of
-                             Just line -> putStrLn $ "\ESC[32m" ++ nodeName (procNode p) ++ "+> " ++ line ++ "\ESC[0m"
-                             Nothing -> putStrLn $ "\ESC[31m" ++ nodeName (procNode p) ++ "/> expect failed" ++ "\ESC[0m"
-                    Left err -> putStrLn $ "failed to parse re: " ++ err
-
-send :: Process -> String -> IO ()
-send p str = do
-    hPutStrLn (procStdin p) str
+expect :: Process -> Regex -> IO ()
+expect p re = do
+    mbmatch <- atomically $ do
+        out <- readTVar (procOutput p)
+        case tryMatch re out of
+             Nothing -> retry
+             Just (m, out') -> do
+                 writeTVar (procOutput p) out'
+                 return $ Just m
+    case mbmatch of
+         Just line -> putStrLn $ "\ESC[32m" ++ unpackNodeName (nodeName (procNode p)) ++ "+> " ++ line ++ "\ESC[0m"
+         Nothing -> putStrLn $ "\ESC[31m" ++ unpackNodeName (nodeName (procNode p)) ++ "/> expect failed" ++ "\ESC[0m"
+
+send :: Process -> Text -> IO ()
+send p line = do
+    T.hPutStrLn (procStdin p) line
     hFlush (procStdin p)
 
-main :: IO ()
-main = do
-    [tool] <- getArgs
-
+runTest :: String -> Test -> IO ()
+runTest tool test = do
     net <- initNetwork
-    node1 <- getNode net 1
-    node2 <- getNode net 2
 
-    p1 <- spawnOn node1 tool
-    p2 <- spawnOn node2 tool
+    forM_ (testSteps test) $ \case
+        Spawn pname nname -> do
+            node <- getNode net nname
+            void $ spawnOn node pname tool
+
+        Send pname line -> do
+            p <- getProcess net pname
+            send p line
 
-    send p1 "create-identity Device1"
-    send p2 "create-identity Device2"
-    send p1 "start-server"
-    send p2 "start-server"
-    expect p1 "peer [0-9]+ 192.168.0.11:29665"
-    expect p1 "peer [0-9]+ 192.168.0.12:29665"
-    expect p2 "peer [0-9]+ 192.168.0.12:29665"
-    expect p2 "peer [0-9]+ 192.168.0.11:29665"
+        Expect pname regex -> do
+            p <- getProcess net pname
+            expect p regex
 
     exitNetwork net
+
+main :: IO ()
+main = do
+    [tool] <- getArgs
+
+    let pat1 = "peer [0-9]+ 192.168.0.11:29665"
+    let pat2 = "peer [0-9]+ 192.168.0.12:29665"
+    Right re1 <- return $ compile defaultCompOpt defaultExecOpt ("^" ++ pat1 ++ "$")
+    Right re2 <- return $ compile defaultCompOpt defaultExecOpt ("^" ++ pat2 ++ "$")
+
+    runTest tool Test
+        { testName = T.pack "Test"
+        , testSteps =
+            [ Spawn (ProcName (T.pack "p1")) (NodeName (T.pack "n1"))
+            , Spawn (ProcName (T.pack "p2")) (NodeName (T.pack "n2"))
+            , Send (ProcName (T.pack "p1")) (T.pack "create-identity Device1")
+            , Send (ProcName (T.pack "p2")) (T.pack "create-identity Device2")
+            , Send (ProcName (T.pack "p1")) (T.pack "start-server")
+            , Send (ProcName (T.pack "p2")) (T.pack "start-server")
+            , Expect (ProcName (T.pack "p1")) re1
+            , Expect (ProcName (T.pack "p1")) re2
+            , Expect (ProcName (T.pack "p2")) re2
+            , Expect (ProcName (T.pack "p2")) re1
+            ]
+        }
diff --git a/src/Test.hs b/src/Test.hs
new file mode 100644
index 0000000..f26e67a
--- /dev/null
+++ b/src/Test.hs
@@ -0,0 +1,33 @@
+module Test (
+    Test(..),
+    TestStep(..),
+
+    ProcName(..), unpackProcName,
+    NodeName(..), unpackNodeName,
+) where
+
+import Data.Text (Text)
+import qualified Data.Text as T
+
+import Text.Regex.TDFA
+
+data Test = Test
+    { testName :: Text
+    , testSteps :: [TestStep]
+    }
+
+data TestStep = Spawn ProcName NodeName
+              | Send ProcName Text
+              | Expect ProcName Regex
+
+newtype ProcName = ProcName Text
+    deriving (Eq, Ord)
+
+unpackProcName :: ProcName -> String
+unpackProcName (ProcName tname) = T.unpack tname
+
+newtype NodeName = NodeName Text
+    deriving (Eq, Ord)
+
+unpackNodeName :: NodeName -> String
+unpackNodeName (NodeName tname) = T.unpack tname