{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
module Main (main) where
import Control.Monad (replicateM_)
import qualified Control.NumerusClosus as NC
import Control.NumerusClosus ((.&&), (.||))
import Data.Either (isLeft, isRight)
import Data.List.NonEmpty (NonEmpty ((:|)))
import qualified Data.List.NonEmpty as NE
import Data.Time (NominalDiffTime, UTCTime (..), addUTCTime, fromGregorian, secondsToNominalDiffTime)
import qualified Hedgehog
import qualified Hedgehog.Gen as Gen
import qualified Hedgehog.Range as Range
import Test.Hspec (Expectation, Spec, describe, hspec, it, shouldBe, shouldSatisfy)
import Test.Hspec.Hedgehog (PropertyT, forAll, hedgehog, (===))
main :: IO ()
main = hspec spec
baseTime :: UTCTime
baseTime = UTCTime (fromGregorian 2024 1 1) 0
offsetTime :: Int -> UTCTime
offsetTime n = addUTCTime (secondsToNominalDiffTime (fromIntegral n)) baseTime
nextTimeUnit :: UTCTime -> UTCTime
nextTimeUnit = addUTCTime 0.000001
genUTCTime :: Hedgehog.Gen UTCTime
genUTCTime = do
s <- Gen.integral (Range.linear 0 86400)
pure $ addUTCTime (secondsToNominalDiffTime (fromIntegral s)) baseTime
debitN :: Int -> UTCTime -> NC.RateLimiter -> Either NC.NextDebitable NC.RateLimiter
debitN 0 _ rl = Right rl
debitN n t rl = case NC.debit rl t of
Right rl' -> debitN (n - 1) t rl'
Left nd -> Left nd
debitSequence :: [UTCTime] -> NC.RateLimiter -> Either NC.NextDebitable NC.RateLimiter
debitSequence [] rl = Right rl
debitSequence (t:ts) rl = case NC.debit rl t of
Right rl' -> debitSequence ts rl'
Left nd -> Left nd
assertRight :: Either NC.NextDebitable NC.RateLimiter -> Expectation
assertRight (Right _) = pure ()
assertRight (Left nd) = fail $ "Expected Right, got Left " ++ show nd
assertLeftNever :: Either NC.NextDebitable NC.RateLimiter -> Expectation
assertLeftNever (Left NC.Never) = pure ()
assertLeftNever (Left nd) = fail $ "Expected Left Never, got Left " ++ show nd
assertLeftNever (Right _) = fail "Expected Left Never, got Right"
assertLeftDebitableFrom :: UTCTime -> Either NC.NextDebitable NC.RateLimiter -> Expectation
assertLeftDebitableFrom t (Left (NC.DebitableFrom t')) | t == t' = pure ()
assertLeftDebitableFrom t (Left nd) = fail $ "Expected Left (DebitableFrom " ++ show t ++ "), got Left " ++ show nd
assertLeftDebitableFrom _ (Right _) = fail "Expected Left DebitableFrom, got Right"
assertLeftAny :: Either NC.NextDebitable NC.RateLimiter -> Expectation
assertLeftAny (Left _) = pure ()
assertLeftAny (Right _) = fail "Expected Left, got Right"
assertLeftNeverH :: (Hedgehog.MonadTest m) => Either NC.NextDebitable NC.RateLimiter -> m ()
assertLeftNeverH (Left NC.Never) = pure ()
assertLeftNeverH _ = Hedgehog.failure
assertLeftDebitableFromH :: (Hedgehog.MonadTest m) => UTCTime -> Either NC.NextDebitable NC.RateLimiter -> m ()
assertLeftDebitableFromH t (Left (NC.DebitableFrom t')) | t == t' = pure ()
assertLeftDebitableFromH _ _ = Hedgehog.failure
spec :: Spec
spec = do
describe "Control.NumerusClosus" do
describe "alwaysAllow" do
it "debiting always returns Right" do
assertRight $ NC.debit NC.alwaysAllow baseTime
assertRight $ debitN 100 baseTime NC.alwaysAllow
describe "alwaysDeny" do
it "debiting always returns Left Never" do
assertLeftNever $ NC.debit NC.alwaysDeny baseTime
describe "finiteBucket" do
it "Allows exactly n debits, then denies with Never" do
let rl = NC.finiteBucket 3
case debitN 3 baseTime rl of
Right rl' -> assertLeftNever $ NC.debit rl' baseTime
Left _ -> fail "should allow 3 debits"
it "finiteBucket 0 denies immediately" do
assertLeftNever $ NC.debit (NC.finiteBucket 0) baseTime
it "finiteBucket 1 allows once then denies" do
let rl = NC.finiteBucket 1
case NC.debit rl baseTime of
Right rl' -> assertLeftNever $ NC.debit rl' baseTime
Left _ -> fail "should allow 1 debit"
describe "fixedWindow" do
let windowSize = 60
maxBucket = 5 :: Int
rl = NC.fixedWindow windowSize (fromIntegral maxBucket) baseTime
it "Allows maxBucket debits within window" do
assertRight $ debitN maxBucket baseTime rl
it "Denies after maxBucket debits within same window" do
case debitN maxBucket baseTime rl of
Right rl' -> assertLeftDebitableFrom (nextTimeUnit (offsetTime 60)) $ NC.debit rl' baseTime
Left _ -> fail "should allow maxBucket debits"
it "Resets after window expires" do
case debitN maxBucket baseTime rl of
Right rl' -> assertRight $ NC.debit rl' (offsetTime 61)
Left _ -> fail "should allow maxBucket debits"
it "Returns DebitableFrom with correct next window time" do
case debitN maxBucket baseTime rl of
Right rl' -> assertLeftDebitableFrom (nextTimeUnit (offsetTime 60)) $ NC.debit rl' (offsetTime 10)
Left _ -> fail "should allow maxBucket debits"
describe "slidingWindow" do
let windowSize = 60
maxBucket = 3 :: Int
rl = NC.slidingWindow windowSize (fromIntegral maxBucket)
it "Allows maxBucket debits within window" do
assertRight $ debitSequence [offsetTime 1, offsetTime 2, offsetTime 3] rl
it "Denies when full within window" do
case debitSequence [offsetTime 1, offsetTime 2, offsetTime 3] rl of
Right rl' -> assertLeftDebitableFrom (nextTimeUnit (offsetTime 61)) $ NC.debit rl' (offsetTime 10)
Left _ -> fail "should allow maxBucket debits"
it "Allows again after oldest entry expires from window" do
case debitSequence [offsetTime 1, offsetTime 2, offsetTime 3] rl of
Right rl3 -> do
assertLeftDebitableFrom (nextTimeUnit (offsetTime 61)) $ NC.debit rl3 (offsetTime 30)
assertRight $ NC.debit rl3 (offsetTime 61)
Left _ -> fail "debits failed"
describe "slidingWindowCount" do
let windowSize = 60
maxBucket = 10 :: Int
rl = NC.slidingWindowCount windowSize (fromIntegral maxBucket) baseTime
it "Allows within budget" do
assertRight $ debitN maxBucket baseTime rl
it "Denies when weighted estimate exceeds budget" do
case debitN maxBucket baseTime rl of
Right rl' -> assertLeftAny $ NC.debit rl' baseTime
Left _ -> fail "should allow maxBucket debits"
it "Resets after full window elapses" do
case debitN maxBucket baseTime rl of
Right rl' -> assertRight $ NC.debit rl' (nextTimeUnit (offsetTime 60))
Left _ -> fail "should allow maxBucket debits"
describe "slidingWindowBucketed" do
let windowSize = 60
windowCount = 6
maxBucket = 10 :: Int
rl = NC.slidingWindowBucketed windowSize windowCount (fromIntegral maxBucket)
it "Allows within budget" do
assertRight $ debitN maxBucket baseTime rl
it "Denies when sum exceeds budget" do
case debitN maxBucket baseTime rl of
Right rl' -> assertLeftAny $ NC.debit rl' baseTime
Left _ -> fail "should allow maxBucket debits"
describe "(.&&) combinator" do
let rlAllow = NC.alwaysAllow
rlDeny = NC.alwaysDeny
rlFixed = NC.fixedWindow 60 1 baseTime
it "Both allow -> allows" do
assertRight $ NC.debit (rlAllow .&& rlAllow) baseTime
it "One denies -> denies" do
assertLeftNever $ NC.debit (rlAllow .&& rlDeny) baseTime
assertLeftNever $ NC.debit (rlDeny .&& rlAllow) baseTime
it "Both deny -> merges NextDebitable with max (semigroup)" do
case NC.debit rlFixed baseTime of
Right rlFixed' -> do
let rl1 = rlFixed'
rl2 = NC.alwaysDeny
assertLeftNever $ NC.debit (rl1 .&& rl2) baseTime
Left _ -> fail "should allow 1 debit"
describe "(.||) combinator" do
let rlAllow = NC.alwaysAllow
rlDeny = NC.alwaysDeny
rlFixed = NC.fixedWindow 60 1 baseTime
it "Both allow -> allows" do
assertRight $ NC.debit (rlAllow .|| rlAllow) baseTime
it "One allows -> allows" do
assertRight $ NC.debit (rlAllow .|| rlDeny) baseTime
assertRight $ NC.debit (rlDeny .|| rlAllow) baseTime
it "Both deny -> merges with min of DebitableFrom times" do
case NC.debit rlFixed baseTime of
Right rlFixed' -> do
let rlFixed2 = NC.fixedWindow 120 1 baseTime
case NC.debit rlFixed2 baseTime of
Right rlFixed2' ->
assertLeftDebitableFrom (nextTimeUnit (offsetTime 60)) $ NC.debit (rlFixed' .|| rlFixed2') baseTime
Left _ -> fail "should allow 1 debit"
Left _ -> fail "should allow 1 debit"
it "Both Never -> Never" do
assertLeftNever $ NC.debit (rlDeny .|| rlDeny) baseTime
describe "allOf / anyOf" do
let rlAllow = NC.alwaysAllow
rlDeny = NC.alwaysDeny
it "allOf works like folded .&&" do
assertRight $ NC.debit (NC.allOf (rlAllow :| [rlAllow])) baseTime
assertLeftNever $ NC.debit (NC.allOf (rlAllow :| [rlDeny])) baseTime
it "anyOf works like folded .||" do
assertLeftNever $ NC.debit (NC.anyOf (rlDeny :| [rlDeny])) baseTime
assertRight $ NC.debit (NC.anyOf (rlDeny :| [rlAllow])) baseTime
describe "NextDebitable Semigroup" do
it "Never <> anything = Never" do
NC.Never <> NC.DebitableFrom baseTime `shouldBe` NC.Never
NC.DebitableFrom baseTime <> NC.Never `shouldBe` NC.Never
it "DebitableFrom x <> DebitableFrom y = DebitableFrom (max x y)" do
let t1 = offsetTime 10
t2 = offsetTime 20
NC.DebitableFrom t1 <> NC.DebitableFrom t2 `shouldBe` NC.DebitableFrom t2
describe "scheduleWith" do
it "mock PositionInTime pure test" do
let mockPit = NC.PositionInTime {NC.getTime = pure baseTime, NC.delayUntil = \_ -> pure ()}
rl = NC.finiteBucket 1
x <- NC.scheduleWith mockPit rl (pure True)
case x of
Right (_, v) -> v `shouldBe` True
Left _ -> fail "should have worked"
describe "Properties" do
it "finiteBucket allows exactly n then denies" $ hedgehog do
n <- forAll (Gen.integral (Range.linear 0 100))
let rl = NC.finiteBucket (NC.BucketSize n)
case debitN (fromIntegral n) baseTime rl of
Right rl' -> assertLeftNeverH $ NC.debit rl' baseTime
Left _ -> Hedgehog.failure
it "alwaysAllow never denies" $ hedgehog do
n <- forAll (Gen.integral (Range.linear 0 1000))
case debitN (fromIntegral n) baseTime NC.alwaysAllow of
Right _ -> pure ()
Left _ -> Hedgehog.failure
it "NextDebitable semigroup associativity" $ hedgehog do
a <- forAll genNextDebitable
b <- forAll genNextDebitable
c <- forAll genNextDebitable
(a <> b) <> c === a <> (b <> c)
it "NextDebitable semigroup Never absorption" $ hedgehog do
x <- forAll genNextDebitable
NC.Never <> x === NC.Never
x <> NC.Never === NC.Never
it "fixedWindow allows exactly maxBucket per window" $ hedgehog do
maxBucket <- forAll (Gen.integral (Range.linear 1 50))
let rl = NC.fixedWindow 60 (NC.BucketSize maxBucket) baseTime
case debitN (fromIntegral maxBucket) baseTime rl of
Right rl' -> assertLeftDebitableFromH (nextTimeUnit (offsetTime 60)) $ NC.debit rl' baseTime
Left _ -> Hedgehog.failure
it "(.&&) is at least as restrictive as either operand" $ hedgehog do
t <- forAll genUTCTime
let rl1 = NC.finiteBucket 0
rl2 = NC.alwaysAllow
assertLeftNeverH $ NC.debit (rl1 .&& rl2) t
assertLeftNeverH $ NC.debit (rl2 .&& rl1) t
it "(.||) is at least as permissive as either operand" $ hedgehog do
t <- forAll genUTCTime
let rl1 = NC.finiteBucket 0
rl2 = NC.alwaysAllow
case NC.debit (rl1 .|| rl2) t of
Right _ -> pure ()
Left _ -> Hedgehog.failure
case NC.debit (rl2 .|| rl1) t of
Right _ -> pure ()
Left _ -> Hedgehog.failure
genNextDebitable :: Hedgehog.Gen NC.NextDebitable
genNextDebitable =
Gen.choice
[ pure NC.Never
, NC.DebitableFrom <$> genUTCTime
]