{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NoFieldSelectors #-}

module Skeletest.Internal.Capture (
  captureOutputPlugin,
  FixtureCapturedOutput (..),

  -- * CLI flag
  CaptureOutputFlag (..),
) where

import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.IO qualified as Text
import GHC.IO.Handle qualified as IO
import Skeletest.Internal.CLI (
  FlagSpec (..),
  FormatFlag (..),
  IsFlag (..),
  getFlag,
  getFormatFlag,
 )
import Skeletest.Internal.CLI qualified as CLI
import Skeletest.Internal.Fixtures (
  Fixture (..),
  FixtureSkeletestTmpDir (..),
  getFixture,
  noCleanup,
  withCleanup,
 )
import Skeletest.Internal.Hooks qualified as Hooks
import Skeletest.Internal.Spec.Output (BoxSpecContent (..))
import Skeletest.Internal.TestRunner (
  TestResult (..),
  TestResultMessage (..),
 )
import Skeletest.Plugin (Hooks (..), Plugin (..), defaultHooks, defaultPlugin)
import System.Directory (removePathForcibly)
import System.IO qualified as IO
import UnliftIO.Exception (finally)

captureOutputPlugin :: Plugin
captureOutputPlugin :: Plugin
captureOutputPlugin =
  Plugin
defaultPlugin
    { cliFlags = [CLI.flag @CaptureOutputFlag]
    , hooks = captureOutputHooks
    }

captureOutputHooks :: Hooks
captureOutputHooks :: Hooks
captureOutputHooks =
  Hooks
defaultHooks
    { runTest = Hooks.mkHook $ \RunTestHookContext
_ () -> IO TestResult
run ()
inp -> do
        (CapturedOutput
output, TestResult
result) <- IO TestResult -> IO (CapturedOutput, TestResult)
forall a. IO a -> IO (CapturedOutput, a)
withCaptureOutput (() -> IO TestResult
run ()
inp)
        CapturedOutput -> TestResult -> IO TestResult
addCapturedOutput CapturedOutput
output TestResult
result
    }

newtype CaptureOutputFlag = CaptureOutputFlag Bool

instance IsFlag CaptureOutputFlag where
  flagName :: String
flagName = String
"capture-output"
  flagHelp :: String
flagHelp = String
"Whether to capture stdout/stderr: on (default), off"
  flagSpec :: FlagSpec CaptureOutputFlag
flagSpec =
    OptionalFlag
      { default_ :: CaptureOutputFlag
default_ = Bool -> CaptureOutputFlag
CaptureOutputFlag Bool
True
      , parse :: String -> Either String CaptureOutputFlag
parse = \case
          String
"off" -> CaptureOutputFlag -> Either String CaptureOutputFlag
forall a b. b -> Either a b
Right (CaptureOutputFlag -> Either String CaptureOutputFlag)
-> CaptureOutputFlag -> Either String CaptureOutputFlag
forall a b. (a -> b) -> a -> b
$ Bool -> CaptureOutputFlag
CaptureOutputFlag Bool
False
          String
"on" -> CaptureOutputFlag -> Either String CaptureOutputFlag
forall a b. b -> Either a b
Right (CaptureOutputFlag -> Either String CaptureOutputFlag)
-> CaptureOutputFlag -> Either String CaptureOutputFlag
forall a b. (a -> b) -> a -> b
$ Bool -> CaptureOutputFlag
CaptureOutputFlag Bool
True
          String
s -> String -> Either String CaptureOutputFlag
forall a b. a -> Either a b
Left (String -> Either String CaptureOutputFlag)
-> String -> Either String CaptureOutputFlag
forall a b. (a -> b) -> a -> b
$ String
"invalid value: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
s
      }

type CapturedOutput = Maybe (Text, Text)

withCaptureOutput :: IO a -> IO (CapturedOutput, a)
withCaptureOutput :: forall a. IO a -> IO (CapturedOutput, a)
withCaptureOutput IO a
action = do
  CaptureOutputFlag Bool
output <- IO CaptureOutputFlag
forall a (m :: * -> *). (MonadIO m, IsFlag a) => m a
getFlag
  if Bool
output
    then do
      FixtureCapturedOutputHandles
handles <- forall a (m :: * -> *). (Fixture a, MonadIO m) => m a
getFixture @FixtureCapturedOutputHandles
      (Text
stdout, (Text
stderr, a
a)) <- LogHandle -> IO (Text, a) -> IO (Text, (Text, a))
forall {b}. LogHandle -> IO b -> IO (Text, b)
capture FixtureCapturedOutputHandles
handles.stdout (IO (Text, a) -> IO (Text, (Text, a)))
-> (IO a -> IO (Text, a)) -> IO a -> IO (Text, (Text, a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LogHandle -> IO a -> IO (Text, a)
forall {b}. LogHandle -> IO b -> IO (Text, b)
capture FixtureCapturedOutputHandles
handles.stderr (IO a -> IO (Text, (Text, a))) -> IO a -> IO (Text, (Text, a))
forall a b. (a -> b) -> a -> b
$ IO a
action
      (CapturedOutput, a) -> IO (CapturedOutput, a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Text, Text) -> CapturedOutput
forall a. a -> Maybe a
Just (Text
stdout, Text
stderr), a
a)
    else (CapturedOutput
forall a. Maybe a
Nothing,) (a -> (CapturedOutput, a)) -> IO a -> IO (CapturedOutput, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO a
action
 where
  capture :: LogHandle -> IO b -> IO (Text, b)
capture LogHandle
handle IO b
m =
    Handle -> IO (Text, b) -> IO (Text, b)
forall {b}. Handle -> IO b -> IO b
withRestore LogHandle
handle.real (IO (Text, b) -> IO (Text, b)) -> IO (Text, b) -> IO (Text, b)
forall a b. (a -> b) -> a -> b
$ do
      Handle -> IO ()
IO.hFlush LogHandle
handle.real
      Handle -> Handle -> IO ()
IO.hDuplicateTo LogHandle
handle.log LogHandle
handle.real
      b
a <- IO b
m
      Text
out <- LogHandle -> IO Text
getOutput LogHandle
handle
      (Text, b) -> IO (Text, b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text
out, b
a)

  withRestore :: Handle -> IO b -> IO b
withRestore Handle
h IO b
m = do
    BufferMode
buf <- Handle -> IO BufferMode
IO.hGetBuffering Handle
h
    Handle
orig <- Handle -> IO Handle
IO.hDuplicate Handle
h
    IO b
m IO b -> IO () -> IO b
forall (m :: * -> *) a b. MonadUnliftIO m => m a -> m b -> m a
`finally` do
      Handle -> Handle -> IO ()
IO.hDuplicateTo Handle
orig Handle
h
      Handle -> BufferMode -> IO ()
IO.hSetBuffering Handle
h BufferMode
buf
      Handle -> IO ()
IO.hClose Handle
orig

addCapturedOutput :: CapturedOutput -> TestResult -> IO TestResult
addCapturedOutput :: CapturedOutput -> TestResult -> IO TestResult
addCapturedOutput CapturedOutput
mCapturedOutput TestResult
result = do
  FormatFlag
format <- IO FormatFlag
getFormatFlag
  let output :: [BoxSpecContent]
output = [BoxSpecContent]
-> ((Text, Text) -> [BoxSpecContent])
-> CapturedOutput
-> [BoxSpecContent]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] (Text, Text) -> [BoxSpecContent]
renderOutput CapturedOutput
mCapturedOutput
  TestResult -> IO TestResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TestResult -> IO TestResult) -> TestResult -> IO TestResult
forall a b. (a -> b) -> a -> b
$
    if FormatFlag -> [BoxSpecContent] -> Bool
forall {t :: * -> *} {a}. Foldable t => FormatFlag -> t a -> Bool
shouldShowOutput FormatFlag
format [BoxSpecContent]
output
      then TestResult
result{message = addOutput output result.message}
      else TestResult
result
 where
  renderOutput :: (Text, Text) -> [BoxSpecContent]
renderOutput (Text
stdout, Text
stderr) =
    [[BoxSpecContent]] -> [BoxSpecContent]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
      [ Text -> Text -> [BoxSpecContent]
renderSection Text
"Captured stdout" Text
stdout
      , Text -> Text -> [BoxSpecContent]
renderSection Text
"Captured stderr" Text
stderr
      ]
  renderSection :: Text -> Text -> [BoxSpecContent]
renderSection Text
name Text
s =
    if Text -> Bool
Text.null Text
s
      then []
      else [Text -> BoxSpecContent
BoxHeader Text
name, Text -> BoxSpecContent
BoxText (Text -> BoxSpecContent) -> Text -> BoxSpecContent
forall a b. (a -> b) -> a -> b
$ Text -> Text
Text.stripEnd Text
s]

  shouldShowOutput :: FormatFlag -> t a -> Bool
shouldShowOutput FormatFlag
format t a
output
    | t a -> Bool
forall a. t a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null t a
output = Bool
False
    | FormatFlag
format FormatFlag -> FormatFlag -> Bool
forall a. Eq a => a -> a -> Bool
== FormatFlag
FormatFlag_Verbose = Bool
True
    | TestResult
result.status.success = Bool
False
    | Bool
otherwise = Bool
True

  addOutput :: [BoxSpecContent] -> TestResultMessage -> TestResultMessage
addOutput [BoxSpecContent]
output TestResultMessage
resultMessage =
    [BoxSpecContent] -> TestResultMessage
TestResultMessageBox ([BoxSpecContent] -> TestResultMessage)
-> [BoxSpecContent] -> TestResultMessage
forall a b. (a -> b) -> a -> b
$ TestResultMessage -> [BoxSpecContent]
toBoxContents TestResultMessage
resultMessage [BoxSpecContent] -> [BoxSpecContent] -> [BoxSpecContent]
forall a. Semigroup a => a -> a -> a
<> [BoxSpecContent]
output
  toBoxContents :: TestResultMessage -> [BoxSpecContent]
toBoxContents = \case
    TestResultMessage
TestResultMessageNone -> []
    TestResultMessageInline Text
msg -> [Text -> BoxSpecContent
BoxText Text
msg]
    TestResultMessageBox [BoxSpecContent]
box -> [BoxSpecContent]
box

data FixtureCapturedOutputHandles = FixtureCapturedOutputHandles
  { FixtureCapturedOutputHandles -> LogHandle
stdout :: LogHandle
  , FixtureCapturedOutputHandles -> LogHandle
stderr :: LogHandle
  }

data LogHandle = LogHandle
  { LogHandle -> Handle
log :: IO.Handle
  , LogHandle -> Handle
real :: IO.Handle
  }

initHandle ::
  IO.Handle ->
  FilePath ->
  FilePath ->
  IO (LogHandle, IO ())
initHandle :: Handle -> String -> String -> IO (LogHandle, IO ())
initHandle Handle
real String
tmpdir String
file = do
  (String
fp, Handle
h) <- String -> String -> IO (String, Handle)
IO.openTempFile String
tmpdir String
file
  Handle -> BufferMode -> IO ()
IO.hSetBuffering Handle
h BufferMode
IO.LineBuffering
  let handle :: LogHandle
handle = LogHandle{log :: Handle
log = Handle
h, real :: Handle
real = Handle
real}
      cleanup :: IO ()
cleanup = Handle -> IO ()
IO.hClose Handle
h IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> String -> IO ()
removePathForcibly String
fp
  (LogHandle, IO ()) -> IO (LogHandle, IO ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LogHandle
handle, IO ()
cleanup)

instance Fixture FixtureCapturedOutputHandles where
  fixtureAction :: IO (FixtureCapturedOutputHandles, FixtureCleanup)
fixtureAction = do
    FixtureSkeletestTmpDir String
tmpdir <- IO FixtureSkeletestTmpDir
forall a (m :: * -> *). (Fixture a, MonadIO m) => m a
getFixture
    (LogHandle
stdout, IO ()
cleanupStdout) <- Handle -> String -> String -> IO (LogHandle, IO ())
initHandle Handle
IO.stdout String
tmpdir String
"stdout"
    (LogHandle
stderr, IO ()
cleanupStderr) <- Handle -> String -> String -> IO (LogHandle, IO ())
initHandle Handle
IO.stderr String
tmpdir String
"stderr"
    (FixtureCapturedOutputHandles, FixtureCleanup)
-> IO (FixtureCapturedOutputHandles, FixtureCleanup)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((FixtureCapturedOutputHandles, FixtureCleanup)
 -> IO (FixtureCapturedOutputHandles, FixtureCleanup))
-> (IO () -> (FixtureCapturedOutputHandles, FixtureCleanup))
-> IO ()
-> IO (FixtureCapturedOutputHandles, FixtureCleanup)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FixtureCapturedOutputHandles
-> IO () -> (FixtureCapturedOutputHandles, FixtureCleanup)
forall a. a -> IO () -> (a, FixtureCleanup)
withCleanup FixtureCapturedOutputHandles{LogHandle
stdout :: LogHandle
stderr :: LogHandle
stdout :: LogHandle
stderr :: LogHandle
..} (IO () -> IO (FixtureCapturedOutputHandles, FixtureCleanup))
-> IO () -> IO (FixtureCapturedOutputHandles, FixtureCleanup)
forall a b. (a -> b) -> a -> b
$ do
      IO ()
cleanupStdout
      IO ()
cleanupStderr

getOutput :: LogHandle -> IO Text
getOutput :: LogHandle -> IO Text
getOutput = Integer -> LogHandle -> IO Text
readOutputFrom Integer
0

readOutput :: LogHandle -> IO Text
readOutput :: LogHandle -> IO Text
readOutput LogHandle
handle = do
  Integer
pos <- Handle -> IO Integer
IO.hTell LogHandle
handle.log
  Integer -> LogHandle -> IO Text
readOutputFrom Integer
pos LogHandle
handle

readOutputFrom :: Integer -> LogHandle -> IO Text
readOutputFrom :: Integer -> LogHandle -> IO Text
readOutputFrom Integer
n LogHandle
handle = do
  -- Flush buffers
  Handle -> IO ()
IO.hFlush LogHandle
handle.real
  Handle -> IO ()
IO.hFlush LogHandle
handle.log

  -- Force handle to end of file, to refresh from real handle
  Handle -> SeekMode -> Integer -> IO ()
IO.hSeek LogHandle
handle.log SeekMode
IO.SeekFromEnd Integer
0

  Handle -> SeekMode -> Integer -> IO ()
IO.hSeek LogHandle
handle.log SeekMode
IO.AbsoluteSeek Integer
n
  Text -> IO Text
go Text
""
 where
  go :: Text -> IO Text
go Text
acc = do
    Text
out <- Handle -> IO Text
Text.hGetChunk LogHandle
handle.log
    if Text -> Bool
Text.null Text
out
      then Text -> IO Text
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
acc
      else Text -> IO Text
go (Text -> IO Text) -> Text -> IO Text
forall a b. (a -> b) -> a -> b
$! Text
acc Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
out

-- | Fixture for inspecting the captured output.
--
-- Intended to be used with @OverloadedRecordDot@:
--
-- @
-- output <- getFixture @FixtureCapturedOutput
--
-- -- Read all of stdout/stderr so far
-- stdout <- output.getStdout
-- stderr <- output.getStderr
--
-- -- Read everything in stdout/stderr since the last read
-- stdout_chunk <- output.readStdout
-- stderr_chunk <- output.readStderr
-- @
data FixtureCapturedOutput = FixtureCapturedOutput
  { FixtureCapturedOutput -> IO Text
getStdout :: IO Text
  , FixtureCapturedOutput -> IO Text
getStderr :: IO Text
  , FixtureCapturedOutput -> IO Text
readStdout :: IO Text
  , FixtureCapturedOutput -> IO Text
readStderr :: IO Text
  }

instance Fixture FixtureCapturedOutput where
  fixtureAction :: IO (FixtureCapturedOutput, FixtureCleanup)
fixtureAction = do
    FixtureCapturedOutputHandles
handles <- forall a (m :: * -> *). (Fixture a, MonadIO m) => m a
getFixture @FixtureCapturedOutputHandles
    (FixtureCapturedOutput, FixtureCleanup)
-> IO (FixtureCapturedOutput, FixtureCleanup)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((FixtureCapturedOutput, FixtureCleanup)
 -> IO (FixtureCapturedOutput, FixtureCleanup))
-> (FixtureCapturedOutput
    -> (FixtureCapturedOutput, FixtureCleanup))
-> FixtureCapturedOutput
-> IO (FixtureCapturedOutput, FixtureCleanup)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FixtureCapturedOutput -> (FixtureCapturedOutput, FixtureCleanup)
forall a. a -> (a, FixtureCleanup)
noCleanup (FixtureCapturedOutput
 -> IO (FixtureCapturedOutput, FixtureCleanup))
-> FixtureCapturedOutput
-> IO (FixtureCapturedOutput, FixtureCleanup)
forall a b. (a -> b) -> a -> b
$
      FixtureCapturedOutput
        { getStdout :: IO Text
getStdout = LogHandle -> IO Text
getOutput FixtureCapturedOutputHandles
handles.stdout
        , getStderr :: IO Text
getStderr = LogHandle -> IO Text
getOutput FixtureCapturedOutputHandles
handles.stderr
        , readStdout :: IO Text
readStdout = LogHandle -> IO Text
readOutput FixtureCapturedOutputHandles
handles.stdout
        , readStderr :: IO Text
readStderr = LogHandle -> IO Text
readOutput FixtureCapturedOutputHandles
handles.stderr
        }