haxl-2.5.1.1: tests/TestBadDataSource.hs

{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

{-# LANGUAGE OverloadedStrings #-}
module TestBadDataSource (tests) where

import Haxl.Prelude as Haxl
import Prelude()

import Haxl.Core

import Data.IORef
import Test.HUnit
import Control.Exception
import System.Mem

import ExampleDataSource
import BadDataSource

testEnv impl fn = do
  -- Use allocation limits, just to make sure haxl properly behaves and
  -- doesn't reset this internally somewhere.
  -- `go` will disable this
  setAllocationCounter 5000000 -- 5 meg should be enough, uses ~100k atm
  enableAllocationLimit
  exstate <- ExampleDataSource.initGlobalState
  badstate <- BadDataSource.initGlobalState impl
  let st = stateSet exstate $ stateSet (fn badstate) stateEmpty
  initEnv st ()

wombats :: GenHaxl () () Int
wombats = length <$> listWombats 3

wombatsMany :: GenHaxl () () Int
wombatsMany = length <$> listWombats 7

go :: FetchImpl -> Test
go impl = TestCase $ flip finally disableAllocationLimit $ do
  -- test that a failed acquire doesn't fail the other requests
  ref <- newIORef False
  env <- testEnv impl $ \st ->
    st { failAcquire = throwIO (DataSourceError "acquire")
       , failRelease = writeIORef ref True }

  x <- runHaxl env $
        (dataFetch (FailAfter 0) + wombatsMany)
          `Haxl.catch` \DataSourceError{} -> wombats

  assertEqual "badDataSourceTest1" 3 x

  -- We should *not* have called release
  assertEqual "badDataSourceTest2" False =<< readIORef ref

  -- test that a failed dispatch doesn't fail the other requests
  ref <- newIORef False
  env <- testEnv impl $ \st ->
    st { failDispatch = throwIO (DataSourceError "dispatch")
       , failRelease = writeIORef ref True }

  x <- runHaxl env $
        (dataFetch (FailAfter 0) + wombatsMany)
          `Haxl.catch` \DataSourceError{} -> wombats

  assertEqual "badDataSourceTest3" x 3

  -- We *should* have called release
  assertEqual "badDataSourceTest4" True =<< readIORef ref

  -- test that a failed wait is a DataSourceError
  env <- testEnv impl $ \st ->
    st { failWait = throwIO (DataSourceError "wait") }

  x <- runHaxl env $
        (dataFetch (FailAfter 0) + wombatsMany)
          `Haxl.catch` \DataSourceError{} -> wombats

  assertEqual "badDataSourceTest5" x 3

  -- We *should* have called release
  assertEqual "badDataSourceTest6" True =<< readIORef ref

  -- test that a failed release is still a DataSourceError, even
  -- though the request will have completed successfully
  env <- testEnv impl $ \st ->
    st { failRelease = throwIO (DataSourceError "release") }

  let
    -- In background fetches the scheduler might happen to process the data
    -- source result (FetchError in this case) before it processes the exception
    -- from release. So we have to allow both cases.
    isBg = case impl of
      Background -> True
      BackgroundMVar -> True
      _ -> False
    releaseCatcher e
      | Just DataSourceError{} <- fromException e = wombats
      | Just FetchError{} <- fromException e =
          if isBg then wombats else Haxl.throw e
      | otherwise = Haxl.throw e

  x <- runHaxl env $
        (dataFetch (FailAfter 0) + wombatsMany)
          `Haxl.catch` releaseCatcher

  assertEqual "badDataSourceTest7" x 3

  -- test that if we don't throw anything we get the result
  -- (which is a fetch error for this source)
  env <- testEnv impl id
  x <- runHaxl env $
        (dataFetch (FailAfter 0) + wombatsMany)
          `Haxl.catch` \FetchError{} -> wombats

  assertEqual "badDataSourceTest8" x 3




tests = TestList
  [ TestLabel "badDataSourceTest async" (go Async)
  , TestLabel "badDataSourceTest background" (go Background)
  , TestLabel "badDataSourceTest backgroundMVar" (go BackgroundMVar)
  , TestLabel "badDataSourceTest backgroundFetchSeq" (go BackgroundSeq)
  , TestLabel "badDataSourceTest backgroundFetchPar" (go BackgroundPar)
  ]