Haxl/tests/DataSourceDispatchTests.hs
Aditya Buridi 769213ea57 Adds schedulerHint tests for DataSource (#136)
Summary:
Pull Request resolved: https://github.com/facebook/Haxl/pull/136

Adds unit test for schedulerHint in DataSource. Currently tests the TryToBatch and SubmitImmediately modes. Tests fail if
. Batch size is not equal to the no of requests in `TryToBatch` mode
. More than one request is submitted in `SubmitImmediately` mode.

Reviewed By: DylanZA

Differential Revision: D24922114

fbshipit-source-id: f786b93e9e0b7c7c731ede5fc54780d9e1b87432
2020-11-16 08:14:20 -08:00

81 lines
2.3 KiB
Haskell

-- Copyright (c) 2014-present, Facebook, Inc.
-- All rights reserved.
--
-- This source code is distributed under the terms of a BSD license,
-- found in the LICENSE file.
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ApplicativeDo #-}
module DataSourceDispatchTests (tests) where
import Test.HUnit hiding (State)
import Control.Monad
import Haxl.Core
import Data.Hashable
data DataSourceDispatch ty where
GetBatchSize :: Int -> DataSourceDispatch Int
deriving instance Eq (DataSourceDispatch ty)
deriving instance Show (DataSourceDispatch ty)
instance DataSourceName DataSourceDispatch where
dataSourceName _ = "DataSourceDispatch"
instance StateKey DataSourceDispatch where
data State DataSourceDispatch = DataSourceDispatchState
instance ShowP DataSourceDispatch where showp = show
instance Hashable (DataSourceDispatch a) where
hashWithSalt s (GetBatchSize n) = hashWithSalt s n
initDataSource :: IO (State DataSourceDispatch)
initDataSource = return DataSourceDispatchState
instance DataSource UserEnv DataSourceDispatch where
fetch _state _flags _u = SyncFetch $ \bfs -> forM_ bfs (fill $ length bfs)
where
fill :: Int -> BlockedFetch DataSourceDispatch -> IO ()
fill l (BlockedFetch (GetBatchSize _ ) rv) = putResult rv (Right l)
schedulerHint Batching = TryToBatch
schedulerHint NoBatching = SubmitImmediately
data UserEnv = Batching | NoBatching deriving (Eq)
makeTestEnv :: UserEnv -> IO (Env UserEnv ())
makeTestEnv testUsrEnv = do
st <- initDataSource
e <- initEnv (stateSet st stateEmpty) testUsrEnv
return e { flags = (flags e) { report = 2 } }
schedulerTest:: Test
schedulerTest = TestCase $ do
let
fet = do
x <- dataFetch (GetBatchSize 0)
y <- dataFetch (GetBatchSize 1)
return [x,y]
e <- makeTestEnv Batching
r1 :: [Int] <- runHaxl e fet
assertEqual "Failed to create batches for data fetch" [2,2] r1
eNoBatching <- makeTestEnv NoBatching
r2 :: [Int] <- runHaxl eNoBatching fet
assertEqual "Unexpexted batches in SubmitImmediately" [1,1] r2
return ()
tests :: Test
tests = TestList
[ TestLabel "schedulerTest" schedulerTest
]