From 88a013bc30dd05ea45227f5b4492950f8a63c3db Mon Sep 17 00:00:00 2001
From: sogaiu <983021772@users.noreply.github.com>
Date: Wed, 27 Dec 2023 22:30:30 +0900
Subject: [PATCH] Add some tests for randgen

---
 spork/randgen.janet      |   7 ++
 test/suite-randgen.janet | 140 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 147 insertions(+)
 create mode 100644 test/suite-randgen.janet

diff --git a/spork/randgen.janet b/spork/randgen.janet
index 0c02130..a9e101b 100644
--- a/spork/randgen.janet
+++ b/spork/randgen.janet
@@ -8,6 +8,13 @@
 (defdyn *rng* "RNG used to generate random numbers")
 
 (defn- get-rng
+  ``
+  Get RNG.
+
+  If dynamic variable *rng* is non-nil, use that as RNG.
+  Otherwise create a new RNG and set the dynamic variable
+  *rng* to the newly created value.
+  ``
   []
   (def rng (dyn *rng*))
   (if rng rng (setdyn *rng* (math/rng))))
diff --git a/test/suite-randgen.janet b/test/suite-randgen.janet
new file mode 100644
index 0000000..c562408
--- /dev/null
+++ b/test/suite-randgen.janet
@@ -0,0 +1,140 @@
+(use ../spork/test)
+(use ../spork/randgen)
+
+(start-suite)
+
+(assert-docs "/spork/randgen")
+
+(def delta 0.000000000000001)
+
+(assert (do
+          (set-seed 1)
+          (def expected @[0.205397787403126
+                          0.946474426007273
+                          0.292829399753968
+                          0.579518994067899
+                          0.897969128303416
+                          0.218425627853671
+                          0.583403751886178
+                          0.613190880523174])
+          (def actual
+            (map |(do $ (rand-uniform))
+                 (range (length expected))))
+          (and (all |(> delta (math/abs (- $0 $1)))
+                    expected actual)
+               (all |(or (< 0 $ 1) (zero? $))
+                    actual)))
+        "rand-uniform")
+
+(assert (do
+          (set-seed 1)
+          (def low 0)
+          (def hi 11)
+          (def expected @[2 10 3 6 9 2 6 6])
+          (def actual
+            (map |(do $ (rand-int low hi))
+                 (range (length expected))))
+          (and (deep= expected actual)
+               (all |(<= low $ (dec hi))
+                    actual)))
+        "rand-int")
+
+(def animals [:ant :bee :cat :dog :ewe :fox :gnu :hog])
+
+(assert (do
+          (set-seed 1)
+          (def expected @[1 7 2 4 7 1 4 4])
+          (def actual
+            (map |(do $ (rand-index animals))
+                 (range (length expected))))
+          (and (deep= expected actual)
+               (all |(<= 0 $ (dec (length expected)))
+                    actual)))
+        "rand-index")
+
+(assert (do
+          (set-seed 1)
+          (def expected @[:bee :hog :cat :ewe :hog :bee :ewe :ewe])
+          (def actual
+            (map |(do $ (rand-value animals))
+                 (range (length expected))))
+          (and (deep= expected actual)
+               (all |(index-of $ animals)
+                    actual)))
+        "rand-value")
+
+(assert (deep= (weights-to-cdf [2 6 8])
+               @[0.125 0.5 1])
+        "weights-to-cdf")
+
+(def delta2 0.1)
+
+(assert (do
+          (set-seed 1)
+          (def w1 1)
+          (def w2 2)
+          (def weights [w1 w2])
+          (var n-w1-idx 0)
+          (var n-w2-idx 0)
+          (var other false)
+          (loop [_ :range [0 1000]]
+            (def res (rand-cdf (weights-to-cdf weights)))
+            (cond
+              (zero? res) (++ n-w1-idx)
+              (one? res) (++ n-w2-idx)
+              (set other true)))
+          (and (not other)
+               (> delta2 (math/abs (- (/ w2 w1) (/ n-w2-idx n-w1-idx))))))
+        "rand-cdf")
+
+(assert (do
+          (set-seed 2)
+          (def w1 1)
+          (def w2 2)
+          (def weights [w1 w2])
+          (var n-w1-idx 0)
+          (var n-w2-idx 0)
+          (var other false)
+          (loop [_ :range [0 10000]]
+            (def res (rand-weights weights))
+            (cond
+              (zero? res) (++ n-w1-idx)
+              (one? res) (++ n-w2-idx)
+              (set other true)))
+          (and (not other)
+               (> delta2 (math/abs (- (/ w2 w1) (/ n-w2-idx n-w1-idx))))))
+        "rand-weights")
+
+(assert (do
+          (set-seed 1)
+          (def counts @{})
+          (for _ 0 1000
+            (rand-path (put counts :ant (inc (get counts :ant 0)))
+                       (put counts :bee (inc (get counts :bee 0)))
+                       (put counts :cat (inc (get counts :cat 0)))))
+          (deep= counts @{:ant 344 :bee 318 :cat 338}))
+        "rand-path")
+
+(assert (do
+          (set-seed 1)
+          (def counts @{})
+          (for _ 0 1000
+            (rand-cdf-path (weights-to-cdf [1 2 3])
+                           (put counts :ant (inc (get counts :ant 0)))
+                           (put counts :bee (inc (get counts :bee 0)))
+                           (put counts :cat (inc (get counts :cat 0)))))
+          (deep= counts @{:ant 177 :bee 318 :cat 505}))
+        "rand-cdf-path")
+
+(assert (do
+          (set-seed 1)
+          (def counts @{})
+          (for _ 0 1000
+            (rand-weights-path [1 2 3]
+                               (put counts :ant (inc (get counts :ant 0)))
+                               (put counts :bee (inc (get counts :bee 0)))
+                               (put counts :cat (inc (get counts :cat 0)))))
+          (deep= counts @{:ant 177 :bee 318 :cat 505}))
+        "rand-weights-path")
+
+(end-suite)