content: {-
    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 StandaloneDeriving, GADTs, DeriveDataTypeable #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE ExistentialQuantification #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE TypeFamilies #-}
  module DataCacheTest (tests, newResult, takeResult) where
  
  import Haxl.Core.DataCache as DataCache
  import Haxl.Core.Monad
  import Haxl.Core
  
  import Control.Exception
  import Data.Either
  import Data.Hashable
  import Data.Traversable
  import Data.Typeable
  import Prelude hiding (mapM)
  import Test.HUnit
  import Data.IORef
  import qualified Data.Text as Text
  import Unsafe.Coerce
  
  data TestReq a where
     Req :: Int -> TestReq a -- polymorphic result
    deriving Typeable
  
  deriving instance Eq (TestReq a)
  deriving instance Show (TestReq a)
  
  instance Hashable (TestReq a) where
    hashWithSalt salt (Req i) = hashWithSalt salt i
  
  instance DataSource u TestReq where
    fetch = error "no fetch defined"
  
  instance DataSourceName TestReq where
    dataSourceName _ = Text.pack "TestReq"
  
  instance StateKey TestReq where
    data State TestReq = TestReqState
  
  instance ShowP TestReq where showp = show
  
  data CacheableReq x where CacheableInt :: Int -> CacheableReq Int
    deriving Typeable
  deriving instance Eq (CacheableReq x)
  deriving instance Show (CacheableReq x)
  instance Hashable (CacheableReq x) where
    hashWithSalt s (CacheableInt val) = hashWithSalt s (0::Int, val)
  
  
  newResult :: Monoid w => a -> IO (IVar u w a)
  newResult a = newFullIVar (Ok a mempty)
  
  takeResult :: IVar u w a -> IO (ResultVal a w)
  takeResult IVar{ivarRef = ref} = do
    e <- readIORef ref
    case e of
      IVarFull a -> return a
      _ -> error "takeResult"
  
  
  dcSoundnessTest :: Test
  dcSoundnessTest = TestLabel "DataCache soundness" $ TestCase $ do
    m1 <- newResult 1
    m2 <- newResult "hello"
    cache <- emptyDataCache
    DataCache.insert (Req 2 :: TestReq String) m2 cache
    DataCache.insert (Req 1 :: TestReq Int) m1 cache
  
    -- "Req 1" has a result of type Int, so if we try to look it up
    -- with a result of type String, we should get Nothing, not a crash.
    r <- mapM takeResult =<< DataCache.lookup (Req 1) cache
    assertBool "dcSoundness1" $
      case r :: Maybe (ResultVal String ()) of
       Nothing -> True
       _something_else -> False
  
    r <- mapM takeResult =<< DataCache.lookup (Req 1) cache
    assertBool "dcSoundness2" $
      case r :: Maybe (ResultVal Int ()) of
       Just (Ok 1 Nothing) -> True
       _something_else -> False
  
    r <- mapM takeResult =<< DataCache.lookup (Req 2) cache
    assertBool "dcSoundness3" $
      case r :: Maybe (ResultVal String ()) of
        Just (Ok "hello" Nothing) -> True
        _something_else -> False
  
    r <- mapM takeResult =<< DataCache.lookup (Req 2) cache
    assertBool "dcSoundness4" $
      case r :: Maybe (ResultVal Int ()) of
        Nothing -> True
        _something_else -> False
  
  
  dcStrictnessTest :: Test
  dcStrictnessTest = TestLabel "DataCache strictness" $ TestCase $ do
    env <- initEnv stateEmpty () :: IO (Env () ())
    r <- Control.Exception.try $ runHaxl env $
      cachedComputation (Req (error "BOOM")) $ return "OK"
    assertBool "dcStrictnessTest" $
      case r of
        Left (ErrorCall "BOOM") -> True
        _other -> False
  
  dcFallbackTest :: Test
  dcFallbackTest = TestLabel "DataCache fallback" $ TestList
    [ TestLabel "Base" $ TestCase $ do
        env <- mkEnv
        (r,cached) <- runHaxl env (do
                             a <- dataFetch req
                             b <- cacheResult (CacheableInt 1234) (return 99999)
                             return (a,b))
        (Stats stats) <- readIORef (statsRef env)
        assertEqual "fallback still has stats" 1
          (Prelude.length [x | x@FetchStats{} <- stats])
        assertEqual "dcFallbackTest found" 1 r
        assertEqual "dcFallbackTest cached" 1234 cached
    , TestLabel "Exception" $ TestCase $ do
        env <- mkEnv
        rbad <- Control.Exception.try $ runHaxl env (dataFetch reqBad)
        assertBool "dcFallbackTest not found" $
          case rbad of
            Left (ErrorCall "no fetch defined") -> True
            _ -> False
    , TestLabel "Completions" $ TestCase $ do
        -- check applicative still runs as it would have without a fallback
        env <- mkEnv
        let
          fetchA = dataFetch reqEx
          fetchB = tellWrite 7 >> dataFetch req
        (rbad, writes) <- runHaxlWithWrites env $
          Haxl.Core.try $ (,) <$> fetchA <*> fetchB
        fetches <- countFetches env
        assertEqual "dispatched 2 fetches" 2 fetches
        assertBool "exception propogates" $
          case rbad of
            Left (NotFound _) -> True
            _ -> False
        assertEqual "write side effects happen" [7] (flattenWT writes)
    ]
    where
  
      mkEnv = addLookup <$> initEnv (stateSet TestReqState stateEmpty) ()
  
      countFetches env = do
        (Stats stats) <- readIORef (statsRef env)
        let
          c = sum [ fetchBatchSize x
                  | x@FetchStats{} <- stats
                  ]
        return c
  
      addLookup :: Env () (WriteTree Int) -> Env () (WriteTree Int)
      addLookup e = e { dataCacheFetchFallback = Just (DataCacheLookup lookup)
                      , flags = (flags e) { report = profilingReportFlags }
                      }
      lookup
        :: forall req a . Typeable (req a)
        => req a
        -> IO (Maybe (ResultVal a (WriteTree Int)))
      lookup r
        | typeOf r == typeRep (Proxy :: Proxy (TestReq Int)) =
          -- have to coerce on the way out as results are not Typeable
          -- so you better be sure you do it right!
          return $ unsafeCoerce . doReq <$> cast r
        | typeOf r == typeRep (Proxy :: Proxy (CacheableReq Int)) =
            return $ unsafeCoerce . doCache <$> cast r
        | otherwise = return Nothing
  
      doReq :: TestReq Int -> ResultVal Int (WriteTree Int)
      doReq (Req 999) = ThrowHaxl (toException $ NotFound Text.empty) Nothing
      doReq (Req r) = Ok r Nothing
  
      doCache :: CacheableReq Int -> ResultVal Int (WriteTree Int)
      doCache (CacheableInt i) = Ok i Nothing
  
      req :: TestReq Int
      req = Req 1
  
      reqEx :: TestReq Int
      reqEx = Req 999
  
      reqBad :: TestReq String
      reqBad = Req 2
  
  
  -- tests :: Assertion
  tests = TestList [ dcSoundnessTest
                   , dcStrictnessTest
                   , dcFallbackTest
                   ]
  