From f940434fbd873e90124a46bd1386c59e8cee49f7 Mon Sep 17 00:00:00 2001
From: Luke Lau <luke_lau@icloud.com>
Date: Thu, 12 Jul 2018 14:07:25 +0100
Subject: [PATCH] Add waitForDiagnosticsSource

---
 src/Language/Haskell/LSP/Test.hs | 12 ++++++++++++
 test/Test.hs                     | 10 +++++++++-
 test/data/error/Error.hs         |  2 ++
 3 files changed, 23 insertions(+), 1 deletion(-)
 create mode 100644 test/data/error/Error.hs

diff --git a/src/Language/Haskell/LSP/Test.hs b/src/Language/Haskell/LSP/Test.hs
index bc1faa1..3def440 100644
--- a/src/Language/Haskell/LSP/Test.hs
+++ b/src/Language/Haskell/LSP/Test.hs
@@ -52,6 +52,7 @@ module Language.Haskell.LSP.Test
   , getDocumentSymbols
   -- ** Diagnostics
   , waitForDiagnostics
+  , waitForDiagnosticsSource
   , noDiagnostics
   -- ** Commands
   , executeCommand
@@ -308,6 +309,17 @@ waitForDiagnostics = do
   let (List diags) = diagsNot ^. params . LSP.diagnostics
   return diags
 
+waitForDiagnosticsSource :: String -> Session [Diagnostic]
+waitForDiagnosticsSource src = do
+  diags <- waitForDiagnostics
+  let res = filter matches diags
+  if null res
+    then waitForDiagnosticsSource src
+    else return res
+  where
+    matches :: Diagnostic -> Bool
+    matches d = d ^. source == Just (T.pack src)
+
 -- | Expects a 'PublishDiagnosticsNotification' and throws an
 -- 'UnexpectedDiagnosticsException' if there are any diagnostics
 -- returned.
diff --git a/test/Test.hs b/test/Test.hs
index 005018c..eedc887 100644
--- a/test/Test.hs
+++ b/test/Test.hs
@@ -237,8 +237,16 @@ main = hspec $ do
         , mkRange 75 6 75 22
         , mkRange 71 6 71 22
         ]
-    where mkRange sl sc el ec = Range (Position sl sc) (Position el ec)
 
+  describe "waitForDiagnosticsSource" $
+    it "works" $ runSession "hie --lsp" "test/data/error" $ do
+      openDoc "Error.hs" "haskell"
+      [diag] <- waitForDiagnosticsSource "ghcmod"
+      liftIO $ do
+        diag ^. severity `shouldBe` Just DsError
+        diag ^. source `shouldBe` Just "ghcmod"
+
+mkRange sl sc el ec = Range (Position sl sc) (Position el ec)
 
 didChangeCaps :: ClientCapabilities
 didChangeCaps = def { _workspace = Just workspaceCaps }
diff --git a/test/data/error/Error.hs b/test/data/error/Error.hs
new file mode 100644
index 0000000..79c1dd9
--- /dev/null
+++ b/test/data/error/Error.hs
@@ -0,0 +1,2 @@
+main :: IO Int
+main = return "hello"
-- 
2.30.2