feat(websockets_hello): close stream after main

1 files changed, 3 insertions, 3 deletions

diff --git a/test/js/websockets_hello.ml b/test/js/websockets_hello.ml
index d417e62..316ae5e 100644
--- a/test/js/websockets_hello.ml
+++ b/test/js/websockets_hello.ml
@@ -1,4 +1,5 @@
 open Lwt.Syntax
+open Lwt.Infix
 open Js_of_ocaml
 open Flesh
 
@@ -21,8 +22,7 @@ let main (stream, push) localpart password =
   in let+ _auth = Sasl.send_auth_stanza (stream, push)
                     localpart password
                     Sasl.PLAIN
-     in push None;
-        match _auth with
+     in match _auth with
         | Error (NotAuthorized, Some (_, text)) -> print_endline ("Not authorized: " ^ text)
         | Error (MalformedRequest, Some (_, text)) -> print_endline ("Malformed request: " ^ text)
         | Error _ -> print_endline "Error!"
@@ -40,5 +40,5 @@ let () =
     in
     let* stream, push = Stream.start domain
     in Lwt.catch
-         (fun () -> main (stream, push) localpart password)
+         (fun () -> main (stream, push) localpart password >|= (fun () -> push None))
          (fun exn -> push None; Lwt.fail exn)