From 969fdeb09a3391cfccf7e42c5c0b6850b1cecc31 Mon Sep 17 00:00:00 2001
From: David Sancho Moreno <dsnxmoreno@gmail.com>
Date: Tue, 1 Aug 2023 13:01:49 +0200
Subject: [PATCH] It works, wat

---
 packages/reactDom/src/reactDOM.ml             | 32 +++++++++----------
 .../reactDom/test/test_renderToLwtStream.ml   | 15 +++++++--
 2 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/packages/reactDom/src/reactDOM.ml b/packages/reactDom/src/reactDOM.ml
index 8da89b205..a9b802792 100644
--- a/packages/reactDom/src/reactDOM.ml
+++ b/packages/reactDom/src/reactDOM.ml
@@ -176,23 +176,21 @@ let render_to_stream ~context_state element =
         arr |> Array.to_list |> List.map render_inner |> String.concat ""
     | Upper_case_component component -> (
         print_endline "Upper_case_component";
-        let element =
-          try Some (component ()) with
-          | React.Suspend (Any_promise promise) ->
-              print_endline "| React.Suspend (Any_promise promise) ->";
-              context_state.waiting <- context_state.waiting + 1;
-              Lwt.map
-                (fun _ ->
-                  context_state.push (render_inner element);
-                  context_state.waiting <- context_state.waiting - 1;
-                  if context_state.waiting = 0 then context_state.close ()
-                  else ())
-                promise
-              |> Lwt.ignore_result;
-              None
-          | e -> raise e
-        in
-        match element with Some element -> render_inner element | None -> "")
+        match component () with
+        | element -> render_inner element
+        | exception React.Suspend (Any_promise promise) ->
+            print_endline "| React.Suspend (Any_promise promise) ->";
+            context_state.waiting <- context_state.waiting + 1;
+            Lwt.map
+              (fun _ ->
+                context_state.push (render_inner element);
+                context_state.waiting <- context_state.waiting - 1;
+                if context_state.waiting = 0 then context_state.close () else ())
+              promise
+            |> Lwt.ignore_result;
+            (* TODO: What should it return in this case? *)
+            ""
+        | exception exn -> raise exn)
     | Lower_case_element { tag; attributes; _ }
       when Html.is_self_closing_tag tag ->
         Printf.sprintf "<%s%s />" tag (attributes_to_string tag attributes)
diff --git a/packages/reactDom/test/test_renderToLwtStream.ml b/packages/reactDom/test/test_renderToLwtStream.ml
index b9c89a7e8..6c1af177b 100644
--- a/packages/reactDom/test/test_renderToLwtStream.ml
+++ b/packages/reactDom/test/test_renderToLwtStream.ml
@@ -4,8 +4,19 @@ let assert_string left right =
 let assert_list ty left right =
   Alcotest.check (Alcotest.list ty) "should be equal" right left
 
+module Sleep = struct
+  let cached = ref false
+
+  let delay v =
+    if cached.contents then Lwt.return ()
+    else
+      let open Lwt.Infix in
+      cached.contents <- true;
+      Lwt_unix.sleep v >>= fun () -> Lwt.return ()
+end
+
 let make ~delay =
-  let () = React.use (Lwt_unix.sleep delay) in
+  let () = React.use (Sleep.delay delay) in
   React.createElement "div" [||]
     [
       React.createElement "span" [||]
@@ -28,7 +39,7 @@ let test_silly_stream _switch () : unit Lwt.t =
 let suspense_one _switch () : unit Lwt.t =
   let timer = React.Upper_case_component (fun () -> make ~delay:0.1) in
   let stream, _abort = ReactDOM.renderToLwtStream timer in
-  assert_stream stream [ "Hello"; "1." ]
+  assert_stream stream [ ""; "<div><span>Hello0.1</span></div>" ]
 
 let case title fn = Alcotest_lwt.test_case title `Quick fn