Implement a ReaderT monad Type for mocking network api calls

Question

data TestEnv = TestEnv
  { rateLimiter' :: !TokenBucket
  , apiManager :: !Manager
  , apiKey :: !BS.ByteString
  }
type BunnyReaderT m = ReaderT TestEnv m
class MonadIO m => HasBunny m where
  runRequest :: Request -> m (Response BSL.ByteString)
  applyAuth :: Request -> m Request
  fetchAuth :: m BS.ByteString
  applyAuth req = do
    apiKey <- fetchAuth
    return $ req { requestHeaders = ("AccessKey", apiKey) : requestHeaders req }
  fetchAuth = liftIO $ BS.pack <$> getEnv "AccessKey"
Instance MonadIO m => HasBunny (BunnyReaderT m) where
  runRequest req = do
    config <- ask
    authReq <- applyAuth req
    let burstSize = 75
        toInvRate r = round(1e6/r)
        invRate = toInvRate 75
    liftIO $ tokenBucketWait (rateLimiter' config) burstSize invRate
    liftIO $ httpLbs authReq (apiManager config) 
  fetchAuth = do
    config <- ask
    return $ apiKey config
type TestM = ReaderT TestEnv IO
Instance HasBunny TestM where
... -- to be defined

This is my code snippet for implementing a readerT monad that implements HasBunny typeclass in which the runRequest function can handle parallel api calls with rate limiting.(need reviews if proper rate limiting is applied or not it can only handle 75 request per second)

How do I define the another ReaderT class which implements the same type-class for the test suite, so that network calls can be mocked out

Basically having a TestM monad that behaves differently from BunnyReaderT monad and making sure The test passes with following assertion ---

that only 75 requests per second are made even if a total of 750 concurrent requests are made.

I'm stuck on this problem for a while any help or leads would be highly appreciated.

I'm basically stuck need reviews if my token-bucket implementation is right in rate limiting it to 75 calls per second moreover need help in implementing instance for TestM monad

Answer 1

This is a somewhat complicated design problem, but let me work you through a simple example. Note that there are lots of minor design decisions to make along the way. Because this is an SO answer and not a 10-part blog post, I've avoided talking about all the different alternatives, so this answer shows one way to do it, certainly not the only way, and not necessarily the best way.

A Simple Example

For this answer I'm going to consider a much simplified problem. Suppose we have a program that prints "foo" and "bar" with a configurable delay in between:

{-# LANGUAGE GHC2021 #-}

import Control.Concurrent
import Control.Monad.Reader

data Config = Config { delay :: Int }
  deriving (Show)

type M = ReaderT Config IO

fooBar :: M ()
fooBar = do
  liftIO $ putStrLn "foo"
  liftIO . threadDelay =<< asks delay
  liftIO $ putStrLn "bar"

main :: IO ()
main = runReaderT fooBar (Config 5_000_000)

and we'd like to test it, to make sure the correct strings are printed with the correct timing.

Mocking out putStrLn

To start, maybe we only want to mock out the putStrLn. We can define a single monad type class for our application, MonadApp, with a method for the call we want to mock out, renamed to appPutStrLn to avoid a clash with putStrLn from Prelude. Any MonadApp will also need to be a Monad, so that should be a superclass. In addition, for the monadic effects we don't want to mock out (e.g., accessing Config and performing a threadDelay), including them as superclasses results in the least disruptive top-level type signatures when we rewrite our functions to use a general MonadApp monad:

class (Monad m, MonadReader Config m, MonadIO m) => MonadApp m where
  appPutStrLn :: String -> m ()

Rewriting fooBar to use this class, we get:

fooBar :: MonadApp m => m ()
fooBar = do
  appPutStrLn "foo"
  liftIO . threadDelay =<< asks delay
  appPutStrLn "bar"

In order to run fooBar in production mode, we need to define a concrete monad that implements MonadApp with the original implementation of appPutStrLn. For this purpose, we'll use a newtype for our monad:

newtype App a = App { unApp :: ReaderT Config IO a }

When we try to define an instance for it:

instance MonadApp App

we'll get errors about missing instances for Monad, MonadReader Config, and MonadIO. Even though ReaderT Config IO satisfies all these constraints, the newtype doesn't by default. Using the GeneralizedNewtypeDeriving extension, we can derive these automatically:

newtype App a = App { unApp :: ReaderT Config IO a }
  deriving (Functor, Applicative, Monad, MonadReader Config, MonadIO)

which allows us to write:

instance MonadApp App where
  appPutStrLn = liftIO . putStrLn

It's also helpful to define a "runner" for the App monad:

runApp :: App a -> Config -> IO a
runApp m c = runReaderT (unApp m) c

The resulting main function for production is:

mainProd :: IO ()
mainProd = runApp fooBar (Config 5_000_000)

For the test monad, we want to mock out appPutStrLn so it creates a log of what was printed and when, so we can check whether the right things were printed with the right timing. We'll do this using the RWS monad, with a Reader for the Config and a Writer for a test log:

type TestLog = [TestEntry]
data TestEntry
  -- a time (in seconds) and string
  = AppPutStrLn Double String deriving (Show)

The test monad itself is defined using a newtype:

import Control.Monad.RWS

newtype TestApp a = TestApp { unTestApp :: RWST Config TestLog () IO a }
  deriving (Functor, Applicative, Monad, MonadReader Config, MonadIO, MonadWriter TestLog)

an instance that performs the logging:

import GHC.Clock

instance MonadApp TestApp where
  appPutStrLn str = do
    t <- liftIO getMonotonicTime
    tell [AppPutStrLn t str]

and a runner:

runTestApp :: TestApp a -> Config -> IO (a, TestLog)
runTestApp m c = evalRWST (unTestApp m) c ()

The main function for testing can use a shorter delay, to speed the test. It returns the TestLog, which the test scaffolding can inspect to determine if the output was correct.

mainTest :: IO TestLog
mainTest = snd <$> runTestApp fooBar (Config 100_000)  -- delay for 100ms

The full resulting program is:

{-# LANGUAGE GHC2021 #-}

import Control.Concurrent
import Control.Monad.Reader
import Control.Monad.RWS
import GHC.Clock

-- Application monad class

data Config = Config { delay :: Int }
  deriving (Show)

class (Monad m, MonadReader Config m, MonadIO m) => MonadApp m where
  appPutStrLn :: String -> m ()

-- Core application logic

fooBar :: MonadApp m => m ()
fooBar = do
  appPutStrLn "foo"
  liftIO . threadDelay =<< asks delay
  appPutStrLn "bar"

-- Production monad

newtype App a = App { unApp :: ReaderT Config IO a }
  deriving (Functor, Applicative, Monad, MonadReader Config, MonadIO)

instance MonadApp App where
  appPutStrLn = liftIO . putStrLn

runApp :: App a -> Config -> IO a
runApp m c = runReaderT (unApp m) c

mainProd :: IO ()
mainProd = runApp fooBar (Config 5_000_000)

-- Test monad

type TestLog = [TestEntry]
data TestEntry
  -- a time (in seconds) and string
  = AppPutStrLn Double String deriving (Show)

newtype TestApp a = TestApp { unTestApp :: RWST Config TestLog () IO a }
  deriving (Functor, Applicative, Monad, MonadReader Config, MonadIO, MonadWriter TestLog)

instance MonadApp TestApp where
  appPutStrLn str = do
    t <- liftIO getMonotonicTime
    tell [AppPutStrLn t str]

runTestApp :: TestApp a -> Config -> IO (a, TestLog)
runTestApp m c = evalRWST (unTestApp m) c ()

mainTest :: IO TestLog
mainTest = snd <$> runTestApp fooBar (Config 100_000)  -- delay for 100ms

and running it in test and production mode yields:

λ> mainTest  -- runs in a fraction of a second, prints no output, returns the "log"
[AppPutStrLn 23816.296046313 "foo",AppPutStrLn 23816.396209595 "bar"]
λ> mainProd  -- delays 5 seconds between printing "foo" and "bar"
foo
bar

Using a Different Reader Context

You also asked whether you needed to use the same context (TestEnv, in your example) for both the production and test monads. No, you don't. If, for example, you wanted to add some testing-specific configuration, like a flag indicating whether the appPutStrLn should actually print its output (in addition to logging it) when running in test mode:

data TestConfig = TestConfig { showOutput :: Bool }
  deriving (Show)

then the way you'd do this is by "mocking out" the asks call to fetch from the Config part of the context:

class (Monad m, MonadIO m) => MonadApp m where
  appPutStrLn :: String -> m ()
  appConfig :: (Config -> a) -> m a

and rewriting fooBar to use appConfig in place of asks:

fooBar :: MonadApp m => m ()
fooBar = do
  appPutStrLn "foo"
  liftIO . threadDelay =<< appConfig delay
  appPutStrLn "bar"

The App monad would only contain a Reader Config, as before, since it doesn't need/use the extra TestConfig context:

-- unchanged from previous definition
newtype App a = App { unApp :: ReaderT Config IO a }
  deriving (Functor, Applicative, Monad, MonadReader Config, MonadIO)

and you'd just need to update its MonadApp instance with an appropriate definition for appConfig:

instance MonadApp App where
  appPutStrLn = liftIO . putStrLn
  appConfig = asks

The TestApp monad, on the other hand, would be modified to read from both Config and TestConfig contexts:

newtype TestApp a = TestApp { unTestApp :: RWST (Config, TestConfig) TestLog () IO a }
  deriving (Functor, Applicative, Monad, MonadReader (Config, TestConfig), MonadIO, MonadWriter TestLog)

with an appropriate definition of appConfig in its instance, plus an updated appPutStrLn definition:

instance MonadApp TestApp where
  appPutStrLn str = do
    t <- liftIO getMonotonicTime
    tell [AppPutStrLn t str]
    noisy <- asks (showOutput . snd)
    when (noisy) $ liftIO $ putStrLn str
  appConfig f = asks (f . fst)

and appropriately updated runner:

runTestApp :: TestApp a -> Config -> TestConfig -> IO (a, TestLog)
runTestApp m c tc = evalRWST (unTestApp m) (c, tc) ()

Now you can run the test quickly and quietly:

mainTest :: IO TestLog
mainTest = snd <$> runTestApp fooBar (Config 100_000) (TestConfig False)  -- delay for 100ms, without printing output

or with realistic output and delays (while still generating a test log)

mainRealisticTest :: IO TestLog
mainRealisticTest = snd <$> runTestApp fooBar (Config 5_000_000) (TestConfig True)

giving:

λ> mainRealisticTest
foo
...five second delay here...
bar
[AppPutStrLn 24941.609278419 "foo",AppPutStrLn 24946.612022085 "bar"]

The complete program with this modification:

{-# LANGUAGE GHC2021 #-}

import Control.Concurrent
import Control.Monad.Reader
import Control.Monad.RWS
import GHC.Clock

-- Application monad class

data Config = Config { delay :: Int }
  deriving (Show)

class (Monad m, MonadIO m) => MonadApp m where
  appPutStrLn :: String -> m ()
  appConfig :: (Config -> a) -> m a

-- Core application logic

fooBar :: MonadApp m => m ()
fooBar = do
  appPutStrLn "foo"
  liftIO . threadDelay =<< appConfig delay
  appPutStrLn "bar"

-- Production monad

newtype App a = App { unApp :: ReaderT Config IO a }
  deriving (Functor, Applicative, Monad, MonadReader Config, MonadIO)

instance MonadApp App where
  appPutStrLn = liftIO . putStrLn
  appConfig = asks

runApp :: App a -> Config -> IO a
runApp m c = runReaderT (unApp m) c

mainProd :: IO ()
mainProd = runApp fooBar (Config 5_000_000)

-- Test monad

data TestConfig = TestConfig { showOutput :: Bool }
  deriving (Show)

type TestLog = [TestEntry]
data TestEntry
  -- a time (in seconds) and string
  = AppPutStrLn Double String deriving (Show)

newtype TestApp a = TestApp { unTestApp :: RWST (Config, TestConfig) TestLog () IO a }
  deriving (Functor, Applicative, Monad, MonadReader (Config, TestConfig), MonadIO, MonadWriter TestLog)

instance MonadApp TestApp where
  appPutStrLn str = do
    t <- liftIO getMonotonicTime
    tell [AppPutStrLn t str]
    noisy <- asks (showOutput . snd)
    when (noisy) $ liftIO $ putStrLn str
  appConfig f = asks (f . fst)

runTestApp :: TestApp a -> Config -> TestConfig -> IO (a, TestLog)
runTestApp m c tc = evalRWST (unTestApp m) (c, tc) ()

mainTest :: IO TestLog
mainTest = snd <$> runTestApp fooBar (Config 100_000) (TestConfig False)  -- delay for 100ms, without printing output

mainRealisticTest :: IO TestLog
mainRealisticTest = snd <$> runTestApp fooBar (Config 5_000_000) (TestConfig True)

Mocking out IO

Finally, you could also consider completely mocking out all the IO, including the threadDelay calls. This would allow you to run a "pure" test that simulates the passage of time, allowing you to run time-based tests much faster, without having to decrease delays and/or relax rate limiting.

The resulting complete program might look something like this:

{-# LANGUAGE GHC2021 #-}

import GHC.Clock
import Control.Concurrent
import Control.Monad.Reader
import Control.Monad.RWS

-- Application monad class

data Config = Config { delay :: Int }
  deriving (Show)

class (Monad m, MonadIO m) => MonadApp m where
  appPutStrLn :: String -> m ()
  appThreadDelay :: Int -> m ()
  appConfig :: (Config -> a) -> m a

-- Core application logic

fooBar :: MonadApp m => m ()
fooBar = do
  appPutStrLn "foo"
  appThreadDelay =<< appConfig delay
  appPutStrLn "bar"

-- Production monad

newtype App a = App { unApp :: ReaderT Config IO a }
  deriving (Functor, Applicative, Monad, MonadReader Config, MonadIO)

instance MonadApp App where
  appPutStrLn = liftIO . putStrLn
  appThreadDelay = liftIO . threadDelay
  appConfig = asks

runApp :: App a -> Config -> IO a
runApp m c = runReaderT (unApp m) c

mainProd :: IO ()
mainProd = runApp fooBar (Config 5_000_000)

-- Test monad

data TestConfig = TestConfig
  { showOutput :: Bool    -- show actual output
  , realTimings :: Bool   -- use realistic (non-simulated) timings
  } deriving (Show)

data TestState = TestState
  { fakeTime :: Double }
  deriving (Show)

type TestLog = [TestEntry]
data TestEntry
  -- a time (in seconds) and string
  = AppPutStrLn Double String deriving (Show)

newtype TestApp a = TestApp { unTestApp :: RWST (Config, TestConfig) TestLog TestState IO a }
  deriving (Functor, Applicative, Monad, MonadIO,
            MonadReader (Config, TestConfig), MonadWriter TestLog, MonadState TestState)

instance MonadApp TestApp where
  appPutStrLn str = do
    TestConfig noisy realtime <- asks snd
    t <- if realtime then liftIO getMonotonicTime
      else do t' <- gets fakeTime
              -- add small time increment
              modify (\s -> s {fakeTime = t' + 0.001})
              return t'
    tell [AppPutStrLn t str]
    when (noisy) $ liftIO $ putStrLn str
  appThreadDelay us = do
    realtime <- asks (realTimings . snd)
    if realtime
      then liftIO (threadDelay us)
      -- simulate passage of time
      else modify (\s -> s {fakeTime = fakeTime s + fromIntegral us / 1_000_000})
  appConfig f = asks (f . fst)

runTestApp :: TestApp a -> Config -> TestConfig -> IO (a, TestLog)
runTestApp m c tc = evalRWST (unTestApp m) (c, tc) (TestState 0)

-- test w/ simulated time
-- (this runs in I/O, but only executes pure code with this particular `TestConfig`)
mainTest :: IO TestLog
mainTest = snd <$> runTestApp fooBar (Config 5_000_000) (TestConfig False False)

-- test w/ real timings
-- (this runs in I/O with actual delays and timings, but no output)
mainTestRealtime :: IO TestLog
mainTestRealtime = snd <$> runTestApp fooBar (Config 100_000) (TestConfig False True)

-- test w/ realistic production output and timings
mainTestRealistic :: IO TestLog
mainTestRealistic = snd <$> runTestApp fooBar (Config 5_000_000) (TestConfig True True)