haxl-2.5.1.1: tests/TestUtils.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 #-}
{-# LANGUAGE CPP #-}
module TestUtils
  ( makeTestEnv
  , expectResultWithEnv
  , expectResult
  , expectFetches
  , testinput
  , id1, id2, id3, id4
  ) where

import TestTypes
import MockTAO
import Haxl.DataSource.ConcurrentIO

import Data.IORef
import Data.Aeson
import Test.HUnit
#if MIN_VERSION_aeson(2,0,0)
import qualified Data.Aeson.KeyMap as KeyMap
#else
import qualified Data.HashMap.Strict as KeyMap
#endif

import Haxl.Core

import Prelude()
import Haxl.Prelude

testinput :: Object
testinput = KeyMap.fromList [
  "A" .= (1 :: Int),
  "B" .= (2 :: Int),
  "C" .= (3 :: Int),
  "D" .= (4 :: Int) ]

id1 :: Haxl Id
id1 = lookupInput "A"

id2 :: Haxl Id
id2 = lookupInput "B"

id3 :: Haxl Id
id3 = lookupInput "C"

id4 :: Haxl Id
id4 = lookupInput "D"

makeTestEnv :: Bool -> IO HaxlEnv
makeTestEnv future = do
  tao <- MockTAO.initGlobalState future
  stio <- mkConcurrentIOState
  let st = stateSet stio $ stateSet tao stateEmpty
  env <- initEnv st testinput
  return env { flags = (flags env) {
    report = setReportFlag ReportFetchStats defaultReportFlags } }

expectResultWithEnv
  :: (Eq a, Show a) => a -> Haxl a -> HaxlEnv -> Assertion
expectResultWithEnv result haxl env = do
  a <- runHaxl env haxl
  assertEqual "result" result a

expectResult :: (Eq a, Show a) => a -> Haxl a -> Bool -> Assertion
expectResult result haxl future = do
  env <- makeTestEnv future
  expectResultWithEnv result haxl env

expectFetches :: (Eq a, Show a) => Int -> Haxl a -> Bool -> Assertion
expectFetches n haxl future = do
  env <- makeTestEnv future
  _ <- runHaxl env haxl
  stats <- readIORef (statsRef env)
  assertEqual "fetches" n (numFetches stats)