module Test.Mockery.Environment (
  withEnvironment
, withModifiedEnvironment
) where

import           Control.Exception.Base
import           Control.Monad
import           System.Environment.Compat

-- |
-- Run the given action within the specified environment.
--
-- Before executing the action, `withEnvironment` backs up the current
-- environment, clears out the environment, and then applies the supplied
-- environment.
-- After the action has completed the original environment is restored.
--
-- __Note__: The environment is global to a process, so tests that modify the
-- environment can no longer be run in parallel.
withEnvironment :: [(String, String)] -> IO a -> IO a
withEnvironment environment action = bracketEnvironment $ do
  setEnvironment environment
  action

-- |
-- Run the given action within an augmented environment.
--
-- Before executing the action, `withModifiedEnvironment` backs up the current
-- environment and then augments it with the supplied values.
-- After the action has completed the original environment is restored.
--
-- __Note__: The environment is global to a process, so tests that modify the
-- environment can no longer be run in parallel.
withModifiedEnvironment :: [(String, String)] -> IO a -> IO a
withModifiedEnvironment environment action = bracketEnvironment $ do
  extendEnvironment environment
  action

bracketEnvironment :: IO a -> IO a
bracketEnvironment = bracket getEnvironment setEnvironment . const

setEnvironment :: [(String, String)] -> IO ()
setEnvironment environment = do
  clearEnvironment
  extendEnvironment environment

extendEnvironment :: [(String, String)] -> IO ()
extendEnvironment environment = forM_ environment $ uncurry setEnv

clearEnvironment :: IO ()
clearEnvironment = do
  environment <- getEnvironment
  forM_ environment (unsetEnv . fst)