From 0cb4c03c9ccccddd195d3088501f7e2da43717cf Mon Sep 17 00:00:00 2001
From: Mateusz Galazyn <mateusz.galazyn@iohk.io>
Date: Thu, 6 Jun 2024 00:19:17 +0200
Subject: [PATCH] Fix flaky test cleanup

---
 cardano-testnet/cardano-testnet.cabal        |  4 +-
 cardano-testnet/src/Testnet/Property/Util.hs | 39 +++++++++++++++++++-
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/cardano-testnet/cardano-testnet.cabal b/cardano-testnet/cardano-testnet.cabal
index 61637535b0d..fed8e29e0c3 100644
--- a/cardano-testnet/cardano-testnet.cabal
+++ b/cardano-testnet/cardano-testnet.cabal
@@ -75,6 +75,7 @@ library
                       , process
                       , random
                       , resourcet
+                      , retry
                       , safe-exceptions
                       , scientific
                       , si-timers
@@ -82,6 +83,7 @@ library
                       , tasty ^>= 1.5
                       , tasty-expected-failure
                       , tasty-hedgehog
+                      , temporary
                       , text
                       , time
                       , transformers
@@ -226,7 +228,7 @@ test-suite cardano-testnet-test
                       , exceptions
                       , filepath
                       , hedgehog
-                      , hedgehog-extras 
+                      , hedgehog-extras
                       , http-conduit
                       , lens-aeson
                       , microlens
diff --git a/cardano-testnet/src/Testnet/Property/Util.hs b/cardano-testnet/src/Testnet/Property/Util.hs
index 9db518d81c7..e2a5d958d25 100644
--- a/cardano-testnet/src/Testnet/Property/Util.hs
+++ b/cardano-testnet/src/Testnet/Property/Util.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE MonoLocalBinds #-}
+{-# LANGUAGE NumericUnderscores #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
@@ -15,14 +16,22 @@ module Testnet.Property.Util
 
 import           Cardano.Api
 
+import           Control.Exception.Safe
+import           Control.Monad
+import qualified Control.Retry as R
 import qualified Data.Aeson as Aeson
 import           GHC.Stack
+import qualified System.Directory as IO
 import qualified System.Environment as IO
+import           System.FilePath ((</>))
 import           System.Info (os)
+import qualified System.IO as IO
+import qualified System.IO.Temp as IO
 import qualified System.IO.Unsafe as IO
 
 import qualified Hedgehog as H
 import qualified Hedgehog.Extras as H
+import qualified Hedgehog.Extras.Stock.CallStack as H
 import           Hedgehog.Internal.Property (MonadTest)
 
 
@@ -43,10 +52,36 @@ integrationRetryWorkspace n workspaceName f = withFrozenCallStack $
   if disableRetries
     then
       integration $
-        H.runFinallies $ H.workspace (workspaceName <> "-no-retries") f
+        H.runFinallies $ workspace (workspaceName <> "-no-retries") f
     else
       integration $ H.retry n $ \i ->
-        H.runFinallies $ H.workspace (workspaceName <> "-" <> show i) f
+        H.runFinallies $ workspace (workspaceName <> "-" <> show i) f
+
+-- | Create a workspace directory which will exist for at least the duration of
+-- the supplied block.
+--
+-- The directory will have the supplied prefix but contain a generated random
+-- suffix to prevent interference between tests
+--
+-- The directory will be deleted if the block succeeds, but left behind if
+-- the block fails.
+-- TODO: this is a version which retries deleting of a workspace on exception - upstream to hedgehog-extras
+workspace :: (MonadTest m, MonadIO m, HasCallStack) => FilePath -> (FilePath -> m ()) -> m ()
+workspace prefixPath f = withFrozenCallStack $ do
+  systemTemp <- H.evalIO IO.getCanonicalTemporaryDirectory
+  maybeKeepWorkspace <- H.evalIO $ IO.lookupEnv "KEEP_WORKSPACE"
+  ws <- H.evalIO $ IO.createTempDirectory systemTemp $ prefixPath <> "-test"
+  H.annotate $ "Workspace: " <> ws
+  H.evalIO $ IO.writeFile (ws </> "module") H.callerModuleName
+  f ws
+  when (os /= "mingw32" && maybeKeepWorkspace /= Just "1") $ do
+    -- try to delete the directory 5 times, 100ms apart
+    let retryPolicy = R.constantDelay 100_000 <> R.limitRetries 10
+        -- retry only on IOExceptions
+        ioExH _ = Handler $ \(_ :: IOException) -> pure True
+    -- For some reason, the temporary directory sometimes fails, lets try multiple times before we fail
+    H.evalIO $ R.recovering retryPolicy [ioExH] . const $
+      IO.removePathForcibly  ws
 
 -- | The 'FilePath' in '(FilePath -> H.Integration ())' is the work space directory.
 -- This is created (and returned) via 'H.workspace'.