From 25cf446cde2d0d14e96c1033938c294b44fc9470 Mon Sep 17 00:00:00 2001
From: thomasjm <tom@codedown.io>
Date: Tue, 16 Jul 2024 02:52:06 -0700
Subject: [PATCH] Be able to obtain ffmpeg on demand

---
 sandwich-webdriver/sandwich-webdriver.cabal   |  1 +
 .../src/Test/Sandwich/WebDriver.hs            | 19 ++++++--
 .../src/Test/Sandwich/WebDriver/Config.hs     |  3 +-
 .../WebDriver/Internal/Binaries/Ffmpeg.hs     | 38 ++++++++++++++-
 .../WebDriver/Internal/Dependencies.hs        |  2 +-
 .../Sandwich/WebDriver/Internal/OnDemand.hs   | 45 ++++++++++++++++++
 .../WebDriver/Internal/StartWebDriver.hs      |  9 ++--
 .../Test/Sandwich/WebDriver/Internal/Types.hs | 18 ++++++++
 .../WebDriver/Internal/Types/Video.hs         |  4 +-
 .../Test/Sandwich/WebDriver/Internal/Video.hs | 46 +++++++++++++------
 .../src/Test/Sandwich/WebDriver/Video.hs      |  7 ++-
 11 files changed, 163 insertions(+), 29 deletions(-)
 create mode 100644 sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/OnDemand.hs

diff --git a/sandwich-webdriver/sandwich-webdriver.cabal b/sandwich-webdriver/sandwich-webdriver.cabal
index 832be7de..e073b05d 100644
--- a/sandwich-webdriver/sandwich-webdriver.cabal
+++ b/sandwich-webdriver/sandwich-webdriver.cabal
@@ -52,6 +52,7 @@ library
       Test.Sandwich.WebDriver.Internal.Capabilities
       Test.Sandwich.WebDriver.Internal.Capabilities.Extra
       Test.Sandwich.WebDriver.Internal.Dependencies
+      Test.Sandwich.WebDriver.Internal.OnDemand
       Test.Sandwich.WebDriver.Internal.Screenshots
       Test.Sandwich.WebDriver.Internal.StartWebDriver
       Test.Sandwich.WebDriver.Internal.Types
diff --git a/sandwich-webdriver/src/Test/Sandwich/WebDriver.hs b/sandwich-webdriver/src/Test/Sandwich/WebDriver.hs
index 6128b8e4..d58aff3e 100644
--- a/sandwich-webdriver/src/Test/Sandwich/WebDriver.hs
+++ b/sandwich-webdriver/src/Test/Sandwich/WebDriver.hs
@@ -83,7 +83,11 @@ introduceWebDriver wdd wdOptions = introduceWebDriver' wdd alloc wdOptions
   where
     alloc wdOptions' = do
       clo <- getSomeCommandLineOptions
-      allocateWebDriver (addCommandLineOptionsToWdOptions clo wdOptions')
+      allocateWebDriver (addCommandLineOptionsToWdOptions clo wdOptions') onDemandOptions
+
+    onDemandOptions = OnDemandOptions {
+      ffmpegToUse = webDriverFfmpeg wdd
+      }
 
 -- | Introduce a 'WebDriver' using the current 'NixContext'.
 -- This will pull everything required from the configured Nixpkgs snapshot.
@@ -102,7 +106,13 @@ introduceWebDriverViaNix wdOptions =
   where
     alloc = do
       clo <- getSomeCommandLineOptions
-      allocateWebDriver (addCommandLineOptionsToWdOptions clo wdOptions)
+
+      nc <- getContext nixContext
+      let onDemandOptions = OnDemandOptions {
+            ffmpegToUse = UseFfmpegFromNixpkgs nc
+            }
+
+      allocateWebDriver (addCommandLineOptionsToWdOptions clo wdOptions) onDemandOptions
 
 -- | Same as 'introduceWebDriver', but with a controllable allocation callback.
 introduceWebDriver' :: forall m context. (
@@ -126,10 +136,11 @@ allocateWebDriver :: (
   )
   -- | Options
   => WdOptions
+  -> OnDemandOptions
   -> ExampleT context m WebDriver
-allocateWebDriver wdOptions = do
+allocateWebDriver wdOptions onDemandOptions = do
   dir <- fromMaybe "/tmp" <$> getCurrentFolder
-  startWebDriver wdOptions dir
+  startWebDriver wdOptions onDemandOptions dir
 
 -- | Clean up the given WebDriver.
 cleanupWebDriver :: (BaseMonad m) => WebDriver -> ExampleT context m ()
diff --git a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Config.hs b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Config.hs
index 54ad675c..3fdd2b94 100644
--- a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Config.hs
+++ b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Config.hs
@@ -11,8 +11,7 @@ module Test.Sandwich.WebDriver.Config (
   , WhenToSave(..)
   , RunMode(..)
 
-  -- * The WebDriver context
-  , WebDriver
+  -- * Accessors for the 'WebDriver' context
   , getWdOptions
   , getDisplayNumber
   , getDownloadDirectory
diff --git a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Binaries/Ffmpeg.hs b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Binaries/Ffmpeg.hs
index 01328330..e1744352 100644
--- a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Binaries/Ffmpeg.hs
+++ b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Binaries/Ffmpeg.hs
@@ -1,7 +1,43 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE TypeApplications #-}
 
 module Test.Sandwich.WebDriver.Internal.Binaries.Ffmpeg (
+  obtainFfmpeg
+
   -- * Types
-  FfmpegToUse(..)
+  , FfmpegToUse(..)
   ) where
 
+import Control.Monad.IO.Unlift
+import Control.Monad.Logger
+import Control.Monad.Reader
+import Data.String.Interpolate
+import qualified Data.Text as T
+import Test.Sandwich
+import Test.Sandwich.Contexts.Files
 import Test.Sandwich.WebDriver.Internal.Binaries.Ffmpeg.Types
+import UnliftIO.Directory
+
+
+-- | Manually obtain an ffmpeg binary, according to the 'FfmpegToUse' policy.
+obtainFfmpeg :: (
+  MonadReader context m, HasBaseContext context
+  , MonadUnliftIO m, MonadLoggerIO m, MonadFail m
+  ) => FfmpegToUse -> m (Either T.Text FilePath)
+obtainFfmpeg UseFfmpegFromPath = findExecutable "ffmpeg" >>= \case
+  Nothing -> return $ Left [i|Couldn't find "ffmpeg" on the PATH.|]
+  Just p -> return $ Right p
+obtainFfmpeg (UseFfmpegAt path) = doesFileExist path >>= \case
+  False -> return $ Left [i|Path '#{path}' didn't exist|]
+  True -> return $ Right path
+obtainFfmpeg (UseFfmpegFromNixpkgs nixContext) =
+  Right <$> getBinaryViaNixDerivation' @"ffmpeg" nixContext ffmpegDerivation
+
+
+ffmpegDerivation :: T.Text
+ffmpegDerivation = [i|
+{ ffmpeg
+}:
+
+ffmpeg.override { withXcb = true; }
+|]
diff --git a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Dependencies.hs b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Dependencies.hs
index f7f85c22..6254f596 100644
--- a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Dependencies.hs
+++ b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Dependencies.hs
@@ -71,7 +71,7 @@ data BrowserDependenciesSpec = BrowserDependenciesSpecChrome {
 -- * Use @firefox@ from the PATH as the browser.
 -- * Download a compatible @geckodriver@ to @\/tmp\/tools@, reusing the one there if found.
 -- * If applicable, it will also get `xvfb-run`, `fluxbox`, and/or `ffmpeg` from the PATH.
-
+--
 -- But, it's easy to customize this behavior. You can define your own 'WebDriverDependencies' and customize
 -- how each of these dependencies are found.
 defaultWebDriverDependencies = WebDriverDependencies {
diff --git a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/OnDemand.hs b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/OnDemand.hs
new file mode 100644
index 00000000..18cb8f3a
--- /dev/null
+++ b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/OnDemand.hs
@@ -0,0 +1,45 @@
+
+module Test.Sandwich.WebDriver.Internal.OnDemand where
+
+import Control.Monad.IO.Unlift
+import Control.Monad.Logger
+import Data.String.Interpolate
+import Data.Text as T
+import Test.Sandwich
+import Test.Sandwich.WebDriver.Internal.Types
+import UnliftIO.Async
+import UnliftIO.Exception
+import UnliftIO.MVar
+
+
+getOnDemand :: forall m a. (
+  MonadUnliftIO m, MonadLogger m
+  ) => MVar (OnDemand a) -> m (Either Text a) -> m a
+getOnDemand onDemandVar doObtain = do
+  result <- modifyMVar onDemandVar $ \case
+    OnDemandErrored msg -> expectationFailure (T.unpack msg)
+    OnDemandNotStarted -> do
+      asy <- async $ do
+        let handler :: SomeException -> m a
+            handler e = do
+              modifyMVar_ onDemandVar (const $ return $ OnDemandErrored [i|Got exception: #{e}|])
+              throwIO e
+
+        handle handler $ do
+          doObtain >>= \case
+            Left err -> do
+              modifyMVar_ onDemandVar (const $ return $ OnDemandErrored err)
+              expectationFailure [i|Failed to obtain: #{err}|]
+
+            Right x -> do
+              modifyMVar_ onDemandVar (const $ return $ OnDemandReady x)
+              return x
+
+      return (OnDemandInProgress asy, Left asy)
+
+    od@(OnDemandInProgress asy) -> pure (od, Left asy)
+    od@(OnDemandReady x) -> pure (od, Right x)
+
+  case result of
+    Right x -> pure x
+    Left asy -> wait asy
diff --git a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/StartWebDriver.hs b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/StartWebDriver.hs
index ba69286a..09526926 100644
--- a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/StartWebDriver.hs
+++ b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/StartWebDriver.hs
@@ -8,7 +8,6 @@
 
 module Test.Sandwich.WebDriver.Internal.StartWebDriver where
 
-import Control.Concurrent
 import Control.Monad
 import Control.Monad.Catch (MonadMask)
 import Control.Monad.IO.Class
@@ -34,6 +33,7 @@ import Test.Sandwich.WebDriver.Internal.Types
 import Test.Sandwich.WebDriver.Internal.Util
 import qualified Test.WebDriver as W
 import UnliftIO.Async
+import UnliftIO.Concurrent
 import UnliftIO.Exception
 import UnliftIO.Process
 import UnliftIO.Timeout
@@ -49,8 +49,8 @@ type Constraints m = (HasCallStack, MonadLogger m, MonadUnliftIO m, MonadMask m)
 startWebDriver :: (
   Constraints m, MonadReader context m
   , HasFile context "java", HasFile context "selenium.jar", HasBrowserDependencies context
-  ) => WdOptions -> FilePath -> m WebDriver
-startWebDriver wdOptions@(WdOptions {capabilities=capabilities'', ..}) runRoot = do
+  ) => WdOptions -> OnDemandOptions -> FilePath -> m WebDriver
+startWebDriver wdOptions@(WdOptions {capabilities=capabilities'', ..}) (OnDemandOptions {..}) runRoot = do
   -- Create a unique name for this webdriver so the folder for its log output doesn't conflict with any others
   webdriverName <- ("webdriver_" <>) <$> liftIO makeUUID
 
@@ -149,6 +149,9 @@ startWebDriver wdOptions@(WdOptions {capabilities=capabilities'', ..}) runRoot =
                           })
             <*> pure downloadDir
 
+            <*> pure ffmpegToUse
+            <*> newMVar OnDemandNotStarted
+
 
 stopWebDriver :: Constraints m => WebDriver -> m ()
 stopWebDriver (WebDriver {wdWebDriver=(h, maybeXvfbSession)}) = do
diff --git a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Types.hs b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Types.hs
index 915b2255..a60c0fe0 100644
--- a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Types.hs
+++ b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Types.hs
@@ -12,12 +12,16 @@ import Data.Default
 import Data.IORef
 import qualified Data.Map as M
 import Data.String.Interpolate
+import Data.Text (Text)
 import Network.HTTP.Client (Manager)
 import System.Process
 import Test.Sandwich
+import Test.Sandwich.WebDriver.Internal.Binaries.Ffmpeg
 import qualified Test.WebDriver as W
 import qualified Test.WebDriver.Class as W
 import qualified Test.WebDriver.Session as W
+import UnliftIO.Async
+
 
 -- | 'Session' is just a 'String' name.
 type Session = String
@@ -64,6 +68,11 @@ data WdOptions = WdOptions {
   -- ^ Number of times to retry an HTTP request if it times out.
   }
 
+data OnDemandOptions = OnDemandOptions {
+  -- | How to obtain ffmpeg binary.
+  ffmpegToUse :: FfmpegToUse
+  }
+
 data HeadlessConfig = HeadlessConfig {
   headlessResolution :: Maybe (Int, Int)
   -- ^ Resolution for the headless browser. Defaults to (1920, 1080)
@@ -96,6 +105,12 @@ defaultWdOptions = WdOptions {
   , httpRetryCount = 0
   }
 
+data OnDemand a =
+  OnDemandNotStarted
+  | OnDemandInProgress (Async a)
+  | OnDemandReady a
+  | OnDemandErrored Text
+
 data WebDriver = WebDriver {
   wdName :: String
   , wdWebDriver :: (ProcessHandle, Maybe XvfbSession)
@@ -103,6 +118,9 @@ data WebDriver = WebDriver {
   , wdSessionMap :: MVar (M.Map Session W.WDSession)
   , wdConfig :: W.WDConfig
   , wdDownloadDir :: FilePath
+
+  , wdFfmpegToUse :: FfmpegToUse
+  , wdFfmpeg :: MVar (OnDemand FilePath)
   }
 
 data InvalidLogsException = InvalidLogsException [W.LogEntry]
diff --git a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Types/Video.hs b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Types/Video.hs
index 1ac561b5..27af549f 100644
--- a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Types/Video.hs
+++ b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Types/Video.hs
@@ -34,7 +34,7 @@ defaultGdigrabOptions :: [String]
 defaultGdigrabOptions = ["-framerate", "30"]
 
 data VideoSettings = VideoSettings {
-  x11grabOptions :: [String]
+  xcbgrabOptions :: [String]
   -- ^ Arguments to x11grab, used with Linux.
   , avfoundationOptions :: [String]
   -- ^ Arguments to avfoundation, used with OS X.
@@ -49,7 +49,7 @@ data VideoSettings = VideoSettings {
 -- | Default video settings.
 defaultVideoSettings :: VideoSettings
 defaultVideoSettings = VideoSettings {
-  x11grabOptions = fastX11VideoOptions
+  xcbgrabOptions = fastX11VideoOptions
   , avfoundationOptions = defaultAvfoundationOptions
   , gdigrabOptions = defaultGdigrabOptions
   , hideMouseWhenRecording = False
diff --git a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Video.hs b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Video.hs
index 332c3e2e..eb851720 100644
--- a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Video.hs
+++ b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Internal/Video.hs
@@ -2,13 +2,21 @@
 
 module Test.Sandwich.WebDriver.Internal.Video where
 
-import Control.Monad.IO.Class
+import Control.Monad.IO.Unlift
+import Control.Monad.Logger
+import Control.Monad.Reader
+import Data.Function
+import qualified Data.List as L
 import Data.Maybe
 import Data.String.Interpolate
-import System.Environment
 import System.Process
+import Test.Sandwich
+import Test.Sandwich.WebDriver.Internal.Binaries.Ffmpeg
+import Test.Sandwich.WebDriver.Internal.OnDemand
 import Test.Sandwich.WebDriver.Internal.Types
 import Test.Sandwich.WebDriver.Internal.Types.Video
+import Test.Sandwich.WebDriver.Types
+import UnliftIO.Environment
 
 #ifdef darwin_HOST_OS
 import Safe
@@ -18,27 +26,38 @@ getMacScreenNumber = undefined
 #endif
 
 
-getVideoArgs :: (MonadIO m) => FilePath -> (Word, Word, Int, Int) -> VideoSettings -> Maybe XvfbSession -> m CreateProcess
+getVideoArgs :: (
+  MonadUnliftIO m, MonadLoggerIO m, MonadFail m
+  , MonadReader context m, HasBaseContext context, HasWebDriverContext context
+  ) => FilePath -> (Word, Word, Int, Int) -> VideoSettings -> Maybe XvfbSession -> m CreateProcess
 getVideoArgs path (width, height, x, y) (VideoSettings {..}) maybeXvfbSession = do
+  WebDriver {wdFfmpeg, wdFfmpegToUse} <- getContext webdriver
+  ffmpeg <- getOnDemand wdFfmpeg (obtainFfmpeg wdFfmpegToUse)
+
 #ifdef linux_HOST_OS
   displayNum <- case maybeXvfbSession of
     Nothing -> fromMaybe "" <$> (liftIO $ lookupEnv "DISPLAY")
     Just (XvfbSession {xvfbDisplayNum}) -> return $ ":" <> show xvfbDisplayNum
 
-  let videoPath = [i|#{path}.avi|]
-  let env' = [("DISPLAY", displayNum)]
+  baseEnv <- getEnvironment
+
   let env = case maybeXvfbSession of
-       Nothing -> Just env'
-       Just (XvfbSession {xvfbXauthority}) -> Just (("XAUTHORITY", xvfbXauthority) : env')
-  let cmd = ["-draw_mouse", (if hideMouseWhenRecording then "0" else "1")
-            , "-y"
+        Nothing -> baseEnv
+        Just (XvfbSession {..}) -> baseEnv
+          & (("DISPLAY", displayNum) :)
+          & (("XAUTHORITY", xvfbXauthority) :)
+          & L.nubBy ((==) `on` fst)
+
+  let videoPath = [i|#{path}.avi|]
+
+  let cmd = ["-y"
             , "-nostdin"
             , "-f", "x11grab"
             , "-s", [i|#{width}x#{height}|]
             , "-i", [i|#{displayNum}.0+#{x},#{y}|]]
-            ++ x11grabOptions
+            ++ xcbgrabOptions
             ++ [videoPath]
-  return ((proc "ffmpeg" cmd) { env = env })
+  return ((proc ffmpeg cmd) { env = Just env })
 #endif
 
 #ifdef darwin_HOST_OS
@@ -52,16 +71,15 @@ getVideoArgs path (width, height, x, y) (VideoSettings {..}) maybeXvfbSession =
                              ++ avfoundationOptions
                              ++ [videoPath]
         Nothing -> error [i|Not launching ffmpeg since OS X screen number couldn't be determined.|]
-  return ((proc "ffmpeg" cmd) { env = Nothing })
+  return ((proc ffmpeg cmd) { env = Nothing })
 #endif
 
 #ifdef mingw32_HOST_OS
   let videoPath = [i|#{path}.mkv|]
   let cmd = ["-f", "gdigrab"
             , "-nostdin"
-            , "-draw_mouse", (if hideMouseWhenRecording then "0" else "1")
             , "-i", "desktop"]
             ++ gdigrabOptions
             ++ [videoPath]
-  return ((proc "ffmpeg.exe" cmd) { env = Nothing })
+  return ((proc ffmpeg cmd) { env = Nothing })
 #endif
diff --git a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Video.hs b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Video.hs
index 7059ceba..268cb80d 100644
--- a/sandwich-webdriver/src/Test/Sandwich/WebDriver/Video.hs
+++ b/sandwich-webdriver/src/Test/Sandwich/WebDriver/Video.hs
@@ -40,7 +40,10 @@ import Test.WebDriver.Commands
 import UnliftIO.Exception
 
 
-type BaseVideoConstraints context m = (MonadLoggerIO m, MonadUnliftIO m, MonadReader context m, HasWebDriverContext context)
+type BaseVideoConstraints context m = (
+  MonadLoggerIO m, MonadUnliftIO m, MonadFail m
+  , MonadReader context m, HasBaseContext context, HasWebDriverContext context
+  )
 
 -- | Wrapper around 'startVideoRecording' which uses the full screen dimensions.
 startFullScreenVideoRecording :: (
@@ -73,7 +76,7 @@ startVideoRecording path (width, height, x, y) vs = do
   sess <- getContext webdriver
   let maybeXvfbSession = getXvfbSession sess
 
-  cp' <- liftIO $ getVideoArgs path (width, height, x, y) vs maybeXvfbSession
+  cp' <- getVideoArgs path (width, height, x, y) vs maybeXvfbSession
   let cp = cp' { create_group = True }
 
   case cmdspec cp of