test: add tests for exit and shutdown notifications (#1241)

* test: add tests for exit and shutdown notifications

Signed-off-by: Rudi Grinberg <me@rgrinberg.com>

Author: jchavarri

ocaml-lsp-server/test/e2e-new/dune

@@ -50,6 +50,7 @@
     completion
     doc_to_md
     document_flow
+    exit_notification
     for_ppx
     hover_extended
     inlay_hints

ocaml-lsp-server/test/e2e-new/exit_notification.ml

@@ -0,0 +1,56 @@
+open Test.Import
+
+let client_capabilities = ClientCapabilities.create ()
+
+module T : sig
+  val run : (unit Client.t -> 'a Fiber.t) -> 'a
+end = struct
+  let run f =
+    let status, a = Test.run_with_status f in
+    let () =
+      match status with
+      | WEXITED n -> Format.eprintf "ocamllsp finished with code = %d@.%!" n
+      | WSIGNALED s -> Format.eprintf "ocamllsp killed with signal = %d@.%!" s
+      | WSTOPPED s -> Format.eprintf "ocamllsp stopped with signal = %d@.%!" s
+    in
+    a
+end
+
+let test run =
+  T.run (fun client ->
+      let run_client () =
+        Client.start
+          client
+          (InitializeParams.create ~capabilities:client_capabilities ())
+      in
+      Fiber.fork_and_join_unit run_client (run client))
+
+let%expect_test "ocamllsp process exits with code 0 after Shutdown and Exit \
+                 notifications are sent" =
+  let run client () =
+    let* (_ : InitializeResult.t) = Client.initialized client in
+    let* () = Client.request client Shutdown in
+    Client.notification client Exit
+  in
+  test run;
+  [%expect {|
+    ocamllsp finished with code = 0  |}]
+
+let%expect_test "ocamllsp does not exit if only Shutdown notification is sent" =
+  let run client () =
+    let* (_ : InitializeResult.t) = Client.initialized client in
+    Client.request client Shutdown
+  in
+  test run;
+  [%expect {|
+    ocamllsp killed with signal = -7  |}]
+
+let%expect_test "ocamllsp process exits with code 0 after Exit notification is \
+                 sent (should be 1)" =
+  let run client () =
+    let* (_ : InitializeResult.t) = Client.initialized client in
+    Client.notification client Exit
+  in
+  test run;
+  [%expect {|
+    ocamllsp finished with code = 0  |}]

ocaml-lsp-server/test/e2e-new/test.ml

@@ -70,6 +70,12 @@ end
 open Import
 
 module T : sig
+  val run_with_status :
+       ?extra_env:string list
+    -> ?handler:unit Client.Handler.t
+    -> (unit Client.t -> 'a Fiber.t)
+    -> Unix.process_status * 'a
+
   val run :
        ?extra_env:string list
     -> ?handler:unit Client.Handler.t
@@ -82,7 +88,7 @@ end = struct
   let bin =
     Bin.which "ocamllsp" ~path:_PATH |> Option.value_exn |> Path.to_string
 
-  let run ?(extra_env = []) ?handler f =
+  let run_with_status ?(extra_env = []) ?handler f =
     let stdin_i, stdin_o = Unix.pipe ~cloexec:true () in
     let stdout_i, stdout_o = Unix.pipe ~cloexec:true () in
     let pid =
@@ -131,26 +137,33 @@ end = struct
       let cancelled = ref false in
       Fiber.fork_and_join_unit
         (fun () ->
-          let+ timeout = Lev_fiber.Timer.Wheel.await timeout in
-          match timeout with
+          Lev_fiber.Timer.Wheel.await timeout >>| function
+          | `Cancelled -> ()
           | `Ok ->
             Unix.kill pid Sys.sigkill;
-            cancelled := true
-          | `Cancelled -> ())
+            cancelled := true)
         (fun () ->
-          let* (_ : Unix.process_status) = Lev_fiber.waitpid ~pid in
-          if !cancelled then Fiber.return ()
-          else Lev_fiber.Timer.Wheel.cancel timeout)
+          let* (server_exit_status : Unix.process_status) =
+            Lev_fiber.waitpid ~pid
+          in
+          let+ () =
+            if !cancelled then Fiber.return ()
+            else Lev_fiber.Timer.Wheel.cancel timeout
+          in
+          server_exit_status)
     in
     Lev_fiber.run (fun () ->
         let* wheel = Lev_fiber.Timer.Wheel.create ~delay:3.0 in
         let+ res = init
-        and+ () =
-          Fiber.all_concurrently_unit
-            [ waitpid wheel; Lev_fiber.Timer.Wheel.run wheel ]
+        and+ status =
+          Fiber.fork_and_join_unit
+            (fun () -> Lev_fiber.Timer.Wheel.run wheel)
+            (fun () -> waitpid wheel)
         in
-        res)
+        (status, res))
     |> Lev_fiber.Error.ok_exn
+
+  let run ?extra_env ?handler f = snd @@ run_with_status ?extra_env ?handler f
 end
 
 include T