fb-util-0.1.0.1: tests/StreamTest.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.
-}

module StreamTest (main) where

import Control.Concurrent
import Data.IORef
import Test.HUnit
import TestRunner

import Facebook.Init
import Util.TimeSec

import Control.Concurrent.Stream
import Data.List

-- Run a worker that sleeps for 1 second over 10 elements.
-- Using 5 workers means we should finish in less than 10 seconds.
timingTest :: Bool -> Test
timingTest bound = TestLabel "timing" $ TestCase $ do
  s <- now
  count <- newIORef (0::Int)
  let
    producer add = mapM_ add [1..(10::Int)]
    sleepy _ = do threadDelay 1000000; atomicModifyIORef' count (\i -> (i+1,()))
  (if bound then streamBound else stream) 5 producer sleepy
  e <- now
  let (TimeSpan d) = timeDiff e s
  assertBool "faster than sequential" (d < 5)
  final <- readIORef count
  assertEqual "all done" 10 final

testMapConcurrently_unordered :: Test
testMapConcurrently_unordered =
  TestLabel "mapConcurrently_unordered" $ TestCase $ do
    let input = [0..9]
    let expected = [1..10]
    startTime <- now
    output <- mapConcurrently_unordered
      10 (\x -> do threadDelay 1000000; pure (x + 1 :: Int)) input
    endTime <- now
    assertBool "Is faster than sequential"
      (toSeconds (timeDiff endTime startTime) < 10)
    assertBool "Received the expected output"
      (expected == sort output)

testMapForConcurrentlyEquivalent_unordered :: Test
testMapForConcurrentlyEquivalent_unordered =
  TestLabel
    "mapConcurrently_unordered==forConcurrently_unordered" $ TestCase $ do
      let input = [0..9]
      mapOut <- mapConcurrently_unordered 10 (\x -> pure $ x + (1 :: Int)) input
      forOut <- forConcurrently_unordered 10 input (\x -> pure $ x + (1 :: Int))
      assertBool
        "mapConcurrently_unordered and forConcurrently_unordered are equivalent"
        (sort mapOut == sort forOut)


main :: IO ()
main = withFacebookUnitTest $ testRunner $ TestList
  [ timingTest True
  , timingTest False
  , testMapConcurrently_unordered
  , testMapForConcurrentlyEquivalent_unordered
  ]