1 files changed, 18 insertions, 5 deletions

diff --git a/server/test/Thermoprint/ServerSpec.hs b/server/test/Thermoprint/ServerSpec.hs
index deb2b9c..028ba2d 100644
--- a/server/test/Thermoprint/ServerSpec.hs
+++ b/server/test/Thermoprint/ServerSpec.hs
@@ -31,6 +31,10 @@ import           Database.Persist.Sqlite
 import           Control.Concurrent
 import           Control.Concurrent.STM
 import           Control.Concurrent.STM.TSem
+import           Control.Concurrent.Async
+
+import           Control.Exception
+import           System.IO.Error
 
 import           System.IO
 import           System.IO.Temp
@@ -63,7 +67,8 @@ data TestManager = TestManager
 
 data RunningServer = RunningServer
   { thread :: ThreadId
-  , termination :: QSem
+  , termination :: MVar (Either SomeException ())
+  , startup :: QSem
   , printer :: TestPrinter
   , manager :: TestManager
   }
@@ -72,7 +77,8 @@ setup :: QSem -> IO RunningServer
 setup startup = withSystemTempFile "thermoprint.sqlite" $ \fp h -> hClose h >> do
   tPrinter <- TestPrinter <$> newEmptyTMVarIO <*> newEmptyTMVarIO
   tManager <- TestManager <$> newEmptyTMVarIO <*> atomically (newTSem 0)
-  termSem <- newQSem 0
+  term <- newEmptyMVar
+  startSem <- newQSem 0
   let
     runSqlite :: ReaderT ConnectionPool (LoggingT IO) a -> IO a
     runSqlite = runStderrLoggingT . filterLogger (const (> LevelDebug)) . withSqlitePool (T.pack fp) 1 . runReaderT
@@ -84,7 +90,7 @@ setup startup = withSystemTempFile "thermoprint.sqlite" $ \fp h -> hClose h >> d
 
     tPM :: MonadIO m => Printout -> m (Maybe PrintingError)
     tPM printout = liftIO . atomically $ putTMVar (outputChan tPrinter) printout >> tryTakeTMVar (failSwitch tPrinter)
-  RunningServer <$> forkFinally (S.thermoprintServer False (Nat runSqlite) $ def' `S.withPrinters` printers) (const $ signalQSem termSem) <*> pure termSem <*> pure tPrinter <*> pure tManager
+  RunningServer <$> forkFinally (S.thermoprintServer False (Nat runSqlite) $ def' `S.withPrinters` printers) (putMVar term) <*> pure term <*> pure startSem <*> pure tPrinter <*> pure tManager
   where
     def' :: MonadIO m => S.Config m
     def' = S.def { S.warpSettings = setBeforeMainLoop (signalQSem startup) $ defaultSettings }
@@ -94,8 +100,15 @@ withSetup = beforeAll setup' . afterAll teardown
   where
     setup' = do
       startup <- newQSem 0
-      setup startup <* waitQSem startup
-    teardown RunningServer{..} = killThread thread >> waitQSem termination
+      r <- setup startup
+      waitStartup <- async $ waitQSem startup
+      waitTermination <- async . readMVar $ termination r
+      wait <- waitEitherCancel waitTermination waitStartup
+      case wait of
+        Left (Left err) -> throwIO err
+        Left (Right _) -> throwIO $ userError "Server thread terminated early"
+        Right _ -> return r
+    teardown RunningServer{..} = killThread thread >> void (readMVar termination)
 
 spec :: Spec
 spec = withSetup $ do