module Test.WebDriver.Commands.ScreenCapture (
  screenshot
  , screenshotElement

  -- * Convenience functions
  , saveScreenshot
  ) where

import Control.Monad.IO.Class
import Data.Aeson as A
import Data.ByteString.Base64.Lazy as B64
import Data.ByteString.Lazy as LBS (ByteString, writeFile)
import qualified Data.Text.Lazy.Encoding as TL
import GHC.Stack
import Test.WebDriver.Types
import Test.WebDriver.Util.Commands


-- | Grab a screenshot of the current page as a PNG image.
screenshot :: (HasCallStack, WebDriver wd) => wd LBS.ByteString
screenshot :: forall (wd :: * -> *).
(HasCallStack, WebDriver wd) =>
wd ByteString
screenshot = (ByteString -> ByteString
B64.decodeLenient (ByteString -> ByteString)
-> (Text -> ByteString) -> Text -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
TL.encodeUtf8) (Text -> ByteString) -> wd Text -> wd ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Method -> Text -> Value -> wd Text
forall (wd :: * -> *) a b.
(HasCallStack, WebDriver wd, ToJSON a, FromJSON b) =>
Method -> Text -> a -> wd b
doSessCommand Method
methodGet Text
"/screenshot" Value
Null

-- | Grab a screenshot of the current page as a PNG image.
screenshotElement :: (HasCallStack, WebDriver wd) => Element -> wd LBS.ByteString
screenshotElement :: forall (wd :: * -> *).
(HasCallStack, WebDriver wd) =>
Element -> wd ByteString
screenshotElement Element
e = (ByteString -> ByteString
B64.decodeLenient (ByteString -> ByteString)
-> (Text -> ByteString) -> Text -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
TL.encodeUtf8) (Text -> ByteString) -> wd Text -> wd ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Method -> Element -> Text -> Value -> wd Text
forall (wd :: * -> *) a b.
(HasCallStack, WebDriver wd, ToJSON a, FromJSON b) =>
Method -> Element -> Text -> a -> wd b
doElemCommand Method
methodGet Element
e Text
"/screenshot" Value
Null

-- | Save a screenshot to a particular location.
saveScreenshot :: (HasCallStack, WebDriver wd) => FilePath -> wd ()
saveScreenshot :: forall (wd :: * -> *).
(HasCallStack, WebDriver wd) =>
FilePath -> wd ()
saveScreenshot FilePath
path = wd ByteString
forall (wd :: * -> *).
(HasCallStack, WebDriver wd) =>
wd ByteString
screenshot wd ByteString -> (ByteString -> wd ()) -> wd ()
forall a b. wd a -> (a -> wd b) -> wd b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO () -> wd ()
forall a. IO a -> wd a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> wd ()) -> (ByteString -> IO ()) -> ByteString -> wd ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FilePath -> ByteString -> IO ()
LBS.writeFile FilePath
path