Split kind tests into component-owned probes

test/RecoveryPhase.hs

@@ -6,6 +6,12 @@ import qualified Data.ByteString as BS
 import qualified Data.ByteString.Lazy as BL
 import           Data.Text (Text)
 import qualified Data.Text as T
+import           Database.PostgreSQL.Simple
+  ( FromRow
+  , Only (..)
+  , Query
+  , ToRow
+  )
 import           Nuthatch.App (newEnv)
 import           Nuthatch.Config
 import           Nuthatch.Db
@@ -16,19 +22,112 @@ import           Nuthatch.Redis
 import           Nuthatch.S3
 import           Nuthatch.Types
 import           Servant.Types.SourceT
-import           System.Environment (getEnv)
+import           System.Environment (getArgs, getEnv)
 import           System.IO (BufferMode (LineBuffering), hSetBuffering, stdout)
+import           Text.Read (readMaybe)
 
 main :: IO ()
 main = do
   hSetBuffering stdout LineBuffering
   cfg <- loadConfig
-  phase <- getEnv "NUTHATCH_RECOVERY_PHASE"
   withDb cfg $ \database -> do
     migrate database "sql" cfg
     withRedis cfg $ \redis -> do
       store <- newS3 cfg
-      runPhase cfg database redis store phase
+      getArgs >>= \case
+        [] -> getEnv "NUTHATCH_RECOVERY_PHASE" >>=
+          runPhase cfg database redis store
+        ["wait-phase", operation, expected] ->
+          waitForPhase database operation expected
+        ["terminal", operation, expected] ->
+          verifyTerminal database operation expected
+        ["active", operation, expectedReservation] ->
+          maybe
+            (fail "invalid expected reservation")
+            (verifyActive database operation)
+            (readMaybe expectedReservation)
+        ["object", expectation, key] ->
+          verifyObject store expectation (ObjectKey (T.pack key))
+        ["cache-invalidated"] ->
+          cacheGet redis "nuthatch:list:recovery-user:shared::" >>=
+            requireEqual "cache invalidation" Nothing
+        ["database-available"] ->
+          waitForQuery database "SELECT 1" ((), Only (1 :: Int))
+        ["log-exhausted"] -> do
+          logBytes <- BS.readFile "/run/nuthatch-recovery/app.log"
+          unless ("\"result\":\"exhausted\"" `BS.isInfixOf` logBytes) $
+            fail "recovery log omitted exhausted retry result"
+        _ -> fail "unknown recovery probe mode"
+
+waitForPhase :: Db -> String -> String -> IO ()
+waitForPhase database operation expected =
+  waitForQuery database
+    "SELECT state || ':' || external_state FROM mutations WHERE operation_id = ?"
+    (Only (T.pack operation), Only (T.pack expected))
+
+verifyTerminal :: Db -> String -> String -> IO ()
+verifyTerminal database operation expected = do
+  waitForQuery database
+    "SELECT state FROM mutations WHERE operation_id = ?"
+    (Only (T.pack operation), Only (T.pack expected))
+  locks <- Db.query database
+    "SELECT count(*) FROM key_locks WHERE operation_id = ?"
+    (Only (T.pack operation))
+    :: IO [Only Int]
+  requireEqual "terminal key lock" [Only 0] locks
+  reservations <- Db.query database
+    "SELECT reserved_bytes FROM pools WHERE pool_id = 'shared'"
+    ()
+    :: IO [Only Integer]
+  requireEqual "terminal reservation" [Only 0] reservations
+
+verifyActive :: Db -> String -> Integer -> IO ()
+verifyActive database operation expectedReservation = do
+  locks <- Db.query database
+    "SELECT count(*) FROM key_locks WHERE operation_id = ?"
+    (Only (T.pack operation))
+    :: IO [Only Int]
+  requireEqual "active key lock" [Only 1] locks
+  reservations <- Db.query database
+    "SELECT reserved_bytes FROM pools WHERE pool_id = 'shared'"
+    ()
+    :: IO [Only Integer]
+  requireEqual "active reservation"
+    [Only expectedReservation]
+    reservations
+
+verifyObject :: S3 -> String -> ObjectKey -> IO ()
+verifyObject store expectation key = do
+  observed <- headObject store key
+  case (expectation, observed) of
+    ("present", Right (Just _)) -> pure ()
+    ("absent", Right Nothing) -> pure ()
+    _ -> fail
+      ("object expectation failed: " <> expectation <> " " <> show observed)
+
+waitForQuery
+  :: (Eq row, Show row, ToRow parameters, FromRow row)
+  => Db
+  -> Query
+  -> (parameters, row)
+  -> IO ()
+waitForQuery database queryText (parameters, expected) =
+  go (120 :: Int)
+  where
+    go remaining = do
+      rows <- Db.query database queryText parameters
+      if rows == [expected]
+        then pure ()
+        else
+          if remaining <= 1
+            then fail
+              ("timed out waiting for database value: " <> show rows)
+            else threadDelay 250000 >> go (remaining - 1)
+
+requireEqual :: (Eq value, Show value) => String -> value -> value -> IO ()
+requireEqual label expected actual =
+  unless (actual == expected) $
+    fail (label <> ": expected " <> show expected <> ", got " <> show actual)
 
 runPhase :: Config -> Db -> Redis -> S3 -> String -> IO ()
 runPhase cfg database redis store phase =