From b3b8d5547a190ab3911e1aa853b27428ab5d6e0f Mon Sep 17 00:00:00 2001
From: Joris Dral <joris@well-typed.com>
Date: Mon, 11 Dec 2023 18:30:37 +0100
Subject: [PATCH] Fix tests

---
 strict-checked-vars/strict-checked-vars.cabal |  9 +++-
 strict-checked-vars/test/Main.hs              | 12 +++--
 .../Class/MonadMVar/Strict/Checked.hs         | 18 ++++---
 .../Class/MonadSTM/Strict/TVar/Checked.hs     | 52 +++++++++++++++++++
 4 files changed, 80 insertions(+), 11 deletions(-)
 create mode 100644 strict-checked-vars/test/Test/Control/Concurrent/Class/MonadSTM/Strict/TVar/Checked.hs

diff --git a/strict-checked-vars/strict-checked-vars.cabal b/strict-checked-vars/strict-checked-vars.cabal
index b7f08c7f2..28903b252 100644
--- a/strict-checked-vars/strict-checked-vars.cabal
+++ b/strict-checked-vars/strict-checked-vars.cabal
@@ -49,7 +49,7 @@ library
 
   default-language: Haskell2010
   build-depends:
-    , base         >=4.9  && <5
+    , base         >=4.9 && <5
     , io-classes   >=1.2 && <1.4
     , strict-mvar  >=1.2 && <1.4
     , strict-stm   >=1.2 && <1.4
@@ -72,6 +72,7 @@ test-suite test
   other-modules:
     Test.Control.Concurrent.Class.MonadMVar.Strict.Checked
     Test.Control.Concurrent.Class.MonadMVar.Strict.Checked.WHNF
+    Test.Control.Concurrent.Class.MonadSTM.Strict.TVar.Checked
     Test.Control.Concurrent.Class.MonadSTM.Strict.TVar.Checked.WHNF
     Test.Utils
 
@@ -90,3 +91,9 @@ test-suite test
     -Wall -Wcompat -Wincomplete-uni-patterns
     -Wincomplete-record-updates -Wpartial-fields -Widentities
     -Wunused-packages -fno-ignore-asserts
+
+  if flag(checkmvarinvariants)
+    cpp-options: -DCHECK_MVAR_INVARIANTS
+
+  if flag(checktvarinvariants)
+    cpp-options: -DCHECK_TVAR_INVARIANTS
\ No newline at end of file
diff --git a/strict-checked-vars/test/Main.hs b/strict-checked-vars/test/Main.hs
index cc170dd78..088f768fb 100644
--- a/strict-checked-vars/test/Main.hs
+++ b/strict-checked-vars/test/Main.hs
@@ -1,11 +1,15 @@
 module Main where
 
-import qualified Test.Control.Concurrent.Class.MonadMVar.Strict.Checked as Test.StrictMVar.Checked
-import qualified Test.Control.Concurrent.Class.MonadSTM.Strict.TVar.Checked.WHNF as Test.StrictTVar.Checked
+import qualified Test.Control.Concurrent.Class.MonadMVar.Strict.Checked
+import qualified Test.Control.Concurrent.Class.MonadMVar.Strict.Checked.WHNF
+import qualified Test.Control.Concurrent.Class.MonadSTM.Strict.TVar.Checked
+import qualified Test.Control.Concurrent.Class.MonadSTM.Strict.TVar.Checked.WHNF
 import           Test.Tasty (defaultMain, testGroup)
 
 main :: IO ()
 main = defaultMain $ testGroup "strict-checked-vars" [
-      Test.StrictMVar.Checked.tests
-    , Test.StrictTVar.Checked.tests
+      Test.Control.Concurrent.Class.MonadMVar.Strict.Checked.tests
+    , Test.Control.Concurrent.Class.MonadMVar.Strict.Checked.WHNF.tests
+    , Test.Control.Concurrent.Class.MonadSTM.Strict.TVar.Checked.tests
+    , Test.Control.Concurrent.Class.MonadSTM.Strict.TVar.Checked.WHNF.tests
     ]
diff --git a/strict-checked-vars/test/Test/Control/Concurrent/Class/MonadMVar/Strict/Checked.hs b/strict-checked-vars/test/Test/Control/Concurrent/Class/MonadMVar/Strict/Checked.hs
index 039746e62..9f6cdd81c 100644
--- a/strict-checked-vars/test/Test/Control/Concurrent/Class/MonadMVar/Strict/Checked.hs
+++ b/strict-checked-vars/test/Test/Control/Concurrent/Class/MonadMVar/Strict/Checked.hs
@@ -1,9 +1,9 @@
+{-# LANGUAGE CPP        #-}
 {-# LANGUAGE RankNTypes #-}
 
 module Test.Control.Concurrent.Class.MonadMVar.Strict.Checked where
 
 import           Control.Concurrent.Class.MonadMVar.Strict.Checked
-import qualified Test.Control.Concurrent.Class.MonadMVar.Strict.Checked.WHNF as Test.WHNF
 import           Test.QuickCheck.Monadic
 import           Test.Tasty
 import           Test.Tasty.QuickCheck
@@ -14,17 +14,16 @@ tests = testGroup "Test.Control.Concurrent.Class.MonadMVar.Strict" [
       testGroup "Checked" [
           testGroup "IO" [
               testProperty "prop_invariantShouldFail" $
-                once $ expectFailure $ monadicIO prop_invariantShouldFail
+                once $ shouldFail $ monadicIO prop_invariantShouldFail
             , testProperty "prop_invariantShouldNotFail" $
-                once                 $ monadicIO prop_invariantShouldNotFail
+                once              $ monadicIO prop_invariantShouldNotFail
             ]
         , testGroup "IOSim" [
               testProperty "prop_invariantShouldFail" $
-                once $ expectFailure $ monadicSim prop_invariantShouldFail
+                once $ shouldFail $ monadicSim prop_invariantShouldFail
             , testProperty "prop_invariantShouldNotFail" $
-                once                 $ monadicSim prop_invariantShouldNotFail
+                once              $ monadicSim prop_invariantShouldNotFail
             ]
-        , Test.WHNF.tests
         ]
     ]
 
@@ -43,3 +42,10 @@ prop_invariantShouldFail :: MonadMVar m => PropertyM m ()
 prop_invariantShouldFail = run $ do
     v <- newMVarWithInvariant invPositiveInt 0
     modifyMVar_ v (\x -> pure $ x - 1)
+
+shouldFail :: Property -> Property
+#if CHECK_MVAR_INVARIANTS
+shouldFail = expectFailure
+#else
+shouldFail = id
+#endif
diff --git a/strict-checked-vars/test/Test/Control/Concurrent/Class/MonadSTM/Strict/TVar/Checked.hs b/strict-checked-vars/test/Test/Control/Concurrent/Class/MonadSTM/Strict/TVar/Checked.hs
new file mode 100644
index 000000000..ef032586d
--- /dev/null
+++ b/strict-checked-vars/test/Test/Control/Concurrent/Class/MonadSTM/Strict/TVar/Checked.hs
@@ -0,0 +1,52 @@
+{-# LANGUAGE CPP        #-}
+{-# LANGUAGE RankNTypes #-}
+
+module Test.Control.Concurrent.Class.MonadSTM.Strict.TVar.Checked where
+
+import           Control.Concurrent.Class.MonadSTM (MonadSTM, atomically)
+import           Control.Concurrent.Class.MonadSTM.Strict.TVar.Checked
+import           Test.QuickCheck.Monadic
+import           Test.Tasty
+import           Test.Tasty.QuickCheck
+import           Test.Utils
+
+tests :: TestTree
+tests = testGroup "Test.Control.Concurrent.Class.MonadSTM.Strict.TVar.Checked" [
+      testGroup "Checked" [
+          testGroup "IO" [
+              testProperty "prop_invariantShouldFail" $
+                once $ shouldFail $ monadicIO prop_invariantShouldFail
+            , testProperty "prop_invariantShouldNotFail" $
+                once              $ monadicIO prop_invariantShouldNotFail
+            ]
+        , testGroup "IOSim" [
+              testProperty "prop_invariantShouldFail" $
+                once $ shouldFail $ monadicSim prop_invariantShouldFail
+            , testProperty "prop_invariantShouldNotFail" $
+                once              $ monadicSim prop_invariantShouldNotFail
+            ]
+        ]
+    ]
+
+-- | Invariant that checks whether an @Int@ is positive.
+invPositiveInt :: Int -> Maybe String
+invPositiveInt x
+  | x >= 0    = Nothing
+  | otherwise = Just $ "x<0 for x=" <> show x
+
+prop_invariantShouldNotFail :: MonadSTM m => PropertyM m ()
+prop_invariantShouldNotFail = run $ atomically $ do
+    v <- newTVarWithInvariant invPositiveInt 0
+    modifyTVar v (+ 1)
+
+prop_invariantShouldFail :: MonadSTM m => PropertyM m ()
+prop_invariantShouldFail = run $ atomically $ do
+    v <- newTVarWithInvariant invPositiveInt 0
+    modifyTVar v (subtract 1)
+
+shouldFail :: Property -> Property
+#if CHECK_TVAR_INVARIANTS
+shouldFail = expectFailure
+#else
+shouldFail = id
+#endif