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.
  -}
  
  module MovingAverageRateLimiterTest (main) where
  
  import Data.List
  import Test.HUnit
  import TestRunner
  
  import Data.MovingAverageRateLimiter
  
  smoothRateLimiter :: Test
  smoothRateLimiter = TestList
    [ TestLabel "Smooth" $ TestCase $ assertEqual "allow 1qps" 0 fails
    , TestLabel "Smooth" $ TestCase $ assertEqual "don't allow 2qps" False
                   (fst $ allow cfg rl' (fromIntegral sz))
    ]
    where
    sz = 50 :: Int
    cfg = mkRateLimiterConfig 1.0 1.0
    rl = mkRateLimiter
    go :: (RateLimiter, Int) -> Int -> (RateLimiter, Int)
    go (r, f) now =
      let (a, r') = allow cfg r (fromIntegral now)
      in (r', f + fromEnum (not a))
    (rl', fails) = foldl' go (rl, 0) [1..sz]
  
  nonSmoothRateLimiter :: Test
  nonSmoothRateLimiter = TestList
    [ TestLabel "NonSmooth" $ TestCase $ assertEqual "allow 2qps" 0 fails
    , TestLabel "NonSmooth" $ TestCase $ assertEqual "don't allow 3qps" False
                   (fst $ allow cfg rl' (fromIntegral sz))
    ]
    where
    sz = 1000 :: Int
    cfg = mkRateLimiterConfig 2.01 10.0
    rl = mkRateLimiter
    go :: (RateLimiter, Int) -> Int -> (RateLimiter, Int)
    go (r, f) now =
      let (a1, r1) = allow cfg r (fromIntegral now)
          (a2, r2) = allow cfg r1 (fromIntegral now)
      in (r2, f + fromEnum (not a1) + fromEnum (not a2))
    (rl', fails) = foldl' go (rl, 0) [1..sz]
  
  observedRateTest :: Test
  observedRateTest = TestList
      [ TestLabel "ObservedRate" $ TestCase $
          assertEqual "don't allow at 2" False a1
      , TestLabel "ObservedRate" $ TestCase $
          assertEqual "allow at 3" True a2
      ]
    where
    cfg = mkRateLimiterConfig 1.0 1.0
    rl = mkRateLimiter
    rl2 = observedRate rl 1.01 2.0
    (a1, rl3) = allow cfg rl2 2
    (a2, _) = allow cfg rl3 3
  
  qpsConverges :: Test
  qpsConverges = TestList
    [ TestLabel "constant" $ TestCase $ convergeTest 1.0 600.0 86400.0
        [ 0.0, 0.1 .. 86400.0 ]
    , TestLabel "skip6hours" $ TestCase $ convergeTest 2.0 60.0 129600.0 $
        [ 0.0, 0.1 .. 21600.0] ++ [ 43200.0, 43200.1 .. 86400 ]
    ]
  
  convergeTest :: Double -> Double -> Double -> [Double] -> Assertion
  convergeTest qps0 hl expected bumps = do
    putStrLn $ "Expected QPS: " ++ show expected
    putStrLn $ "Got: " ++ show count
    assertBool "converges on expected rate" closeEnough
    where
      cfg = mkRateLimiterConfig qps0 hl
      go (r, n) t =
        let (a, r') = allow cfg r t
        in (r', n + fromEnum a)
      (_, count) = foldl' go (mkRateLimiter, 0 :: Int) bumps
      closeEnough = abs (fromIntegral count - expected) / expected < 0.01
  
  main :: IO ()
  main = testRunner $ TestList
    [ smoothRateLimiter
    , nonSmoothRateLimiter
    , observedRateTest
    , qpsConverges
    ]
  