improve waitForDiagnosticsFromSourceWithTimeout

Author: soulomoon

hls-test-utils/src/Test/Hls/Util.hs

@@ -284,31 +284,26 @@ expectNoMoreDiagnostics timeout doc src = do
 -- If timeout is 0 it will wait until the session timeout
 waitForDiagnosticsFromSourceWithTimeout :: Seconds -> TextDocumentIdentifier -> String -> Test.Session [Diagnostic]
 waitForDiagnosticsFromSourceWithTimeout timeout document source = do
-    when (timeout > 0) $
-        -- Give any further diagnostic messages time to arrive.
-        liftIO $ sleep timeout
-        -- Send a dummy message to provoke a response from the server.
-        -- This guarantees that we have at least one message to
-        -- process, so message won't block or timeout.
-    testId <- Test.sendRequest (SMethod_CustomMethod (Proxy @"test")) A.Null
-    handleMessages testId
+    diags <- withTimeout timeout $ waitForDiagnosticsFrom document
+    return $ filter matches diags
   where
     matches :: Diagnostic -> Bool
     matches d = d ^. L.source == Just (T.pack source)
 
-    handleMessages testId = handleDiagnostic testId <|> handleMethod_CustomMethodResponse testId <|> ignoreOthers testId
-    handleDiagnostic testId = do
-        diagsNot <- Test.message SMethod_TextDocumentPublishDiagnostics
-        let fileUri = diagsNot ^. L.params . L.uri
-            diags = diagsNot ^. L.params . L.diagnostics
-            res = filter matches diags
-        if fileUri == document ^. L.uri && not (null res)
-            then return res else handleMessages testId
-    handleMethod_CustomMethodResponse testId = do
-        _ <- Test.responseForId (SMethod_CustomMethod (Proxy @"test")) testId
-        pure []
-
-    ignoreOthers testId = void Test.anyMessage >> handleMessages testId
+withTimeout :: Seconds -> Session a -> Session a
+withTimeout duration f = do
+  chan <- asks messageChan
+  timeoutId <- getCurTimeoutId
+  modify $ \s -> s { overridingTimeout = True }
+  tid <- liftIO $ forkIO $ do
+    threadDelay (ceiling $ duration * 1000000)
+    writeChan chan (TimeoutMessage timeoutId)
+  res <- f
+  liftIO $ killThread tid
+  bumpTimeoutId timeoutId
+  modify $ \s -> s { overridingTimeout = False }
+  return res
+
 
 failIfSessionTimeout :: IO a -> IO a
 failIfSessionTimeout action = action `catch` errorHandler