From 62777e1112f7c625ba6c55fb4c74f808cd22cbdd Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Mon, 19 Jun 2023 15:34:22 -0400
Subject: [PATCH] test: add tests for `Fut.await`

---
 test/await/dune       |  5 ++++
 test/await/t_fib1.ml  | 48 +++++++++++++++++++++++++++++++++++++++
 test/await/t_futs1.ml | 53 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 106 insertions(+)
 create mode 100644 test/await/dune
 create mode 100644 test/await/t_fib1.ml
 create mode 100644 test/await/t_futs1.ml

diff --git a/test/await/dune b/test/await/dune
new file mode 100644
index 00000000..db90ee84
--- /dev/null
+++ b/test/await/dune
@@ -0,0 +1,5 @@
+
+(tests
+  (names t_fib1 t_futs1)
+  (enabled_if (>= %{ocaml_version} 5.0))
+  (libraries moonpool trace tracy-client.trace))
diff --git a/test/await/t_fib1.ml b/test/await/t_fib1.ml
new file mode 100644
index 00000000..99beaac5
--- /dev/null
+++ b/test/await/t_fib1.ml
@@ -0,0 +1,48 @@
+open Moonpool
+
+let rec fib_direct x =
+  if x <= 1 then
+    1
+  else
+    fib_direct (x - 1) + fib_direct (x - 2)
+
+let fib ~on x : int Fut.t =
+  let rec fib_rec x : int =
+    if x <= 18 then
+      fib_direct x
+    else (
+      let t1 = Fut.spawn ~on (fun () -> fib_rec (x - 1))
+      and t2 = Fut.spawn ~on (fun () -> fib_rec (x - 2)) in
+      Fut.await_exn t1 + Fut.await_exn t2
+    )
+  in
+  Fut.spawn ~on (fun () -> fib_rec x)
+
+let () = Tracy_client_trace.setup ()
+let () = assert (List.init 10 fib_direct = [ 1; 1; 2; 3; 5; 8; 13; 21; 34; 55 ])
+
+let fib_40 : int =
+  let pool = Pool.create ~min:8 () in
+  fib ~on:pool 40 |> Fut.wait_block_exn
+
+let () = Printf.printf "fib 40 = %d\n%!" fib_40
+
+let run_test () =
+  let pool = Pool.create ~min:8 () in
+
+  assert (
+    List.init 10 (fib ~on:pool)
+    |> Fut.join_list |> Fut.wait_block_exn
+    = [ 1; 1; 2; 3; 5; 8; 13; 21; 34; 55 ]);
+
+  let fibs = Array.init 3 (fun _ -> fib ~on:pool 40) in
+
+  let res = Fut.join_array fibs |> Fut.wait_block in
+  Pool.shutdown pool;
+
+  assert (res = Ok (Array.make 3 fib_40))
+
+let () =
+  (* now make sure we can do this with multiple pools in parallel *)
+  let jobs = Array.init 2 (fun _ -> Thread.create run_test ()) in
+  Array.iter Thread.join jobs
diff --git a/test/await/t_futs1.ml b/test/await/t_futs1.ml
new file mode 100644
index 00000000..a345925a
--- /dev/null
+++ b/test/await/t_futs1.ml
@@ -0,0 +1,53 @@
+open! Moonpool
+
+let pool = Pool.create ~min:4 ()
+
+let () =
+  let fut = Array.init 10 (fun i -> Fut.spawn ~on:pool (fun () -> i)) in
+  let fut2 = Fut.spawn ~on:pool (fun () -> Array.map Fut.await_exn fut) in
+  assert (Fut.wait_block fut2 = Ok (Array.init 10 (fun x -> x)))
+
+let () =
+  let fut =
+    Array.init 10 (fun i ->
+        Fut.spawn ~on:pool (fun () ->
+            if i < 9 then
+              i
+            else
+              raise Not_found))
+  in
+  let fut2 = Fut.spawn ~on:pool (fun () -> Array.map Fut.await_exn fut) in
+  (* must fail *)
+  assert (Fut.wait_block fut2 |> Result.is_error)
+
+let mk_ret_delay ?(on = pool) n x =
+  Fut.spawn ~on (fun () ->
+      Thread.delay n;
+      x)
+
+let () =
+  let f1 = mk_ret_delay 0.01 1 in
+  let f2 = mk_ret_delay 0.01 2 in
+  let fut = Fut.spawn ~on:pool (fun () -> Fut.await_exn f1, Fut.await_exn f2) in
+  assert (Fut.wait_block_exn fut = (1, 2))
+
+let () =
+  let f1 =
+    let f =
+      Fut.spawn ~on:pool (fun () ->
+          Thread.delay 0.01;
+          1)
+    in
+    Fut.spawn ~on:pool (fun () -> Fut.await_exn f + 1)
+  and f2 =
+    let f =
+      Fut.spawn ~on:pool (fun () ->
+          Thread.delay 0.01;
+          10)
+    in
+    Fut.spawn ~on:pool (fun () ->
+        Thread.delay 0.01;
+        Fut.await_exn f * 2)
+  in
+  let fut = Fut.both f1 f2 in
+  assert (Fut.wait_block fut = Ok (2, 20))