Remove some Lwt stuff from Extra

src/client/remote_nodestream_test.ml

@@ -138,7 +138,7 @@ let set_route_delta port () =
   __wrap__ port conversation
 
 let suite =
-  let w f = Extra.lwt_bracket setup f teardown in
+  let w body = Lwt_extra.OUnit.bracket ~setup ~body ~teardown in
   "nursery" >:::
     ["set_interval" >:: w (set_interval 6666);
      "get_fringe"  >:: w (get_fringe  5555);

src/extra.ml

@@ -22,23 +22,3 @@ let eq_conv conv str i1 i2 =
   let c1 = conv i1 and c2 = conv i2 in
   let msg = Printf.sprintf "%s expected:%s actual:%s" str c1 c2 in
   OUnit.assert_equal ~msg i1 i2
-
-open Lwt
-
-let lwt_bracket setup testcase teardown () =
-  let try_lwt_ f =
-    Lwt.catch f (fun exn -> Lwt.fail exn)
-  in
-  Lwt_main.run
-    begin
-      try_lwt_ setup >>= fun x ->
-      try_lwt_ (fun () ->
-          Lwt.finalize (fun () -> testcase x)
-            (fun () -> teardown x)
-        ) >>= fun () ->
-      Lwt.return ()
-    end
-
-let lwt_test_wrap testcase =
-  let setup = Lwt.return and teardown _ = Lwt.return () in
-  lwt_bracket setup testcase teardown

src/node/catchup_test.ml

@@ -18,7 +18,6 @@ limitations under the License.
 
 open OUnit
 open Lwt
-open Extra
 open Update
 
 let section = Logger.Section.main
@@ -191,7 +190,7 @@ let test_large_tlog_catchup () =
   _tic _fill 60_000 "tcs"
     (fun _store _new_i -> Lwt.return ())
 let suite =
-  let w f = lwt_bracket setup f teardown in
+  let w body = Lwt_extra.OUnit.bracket ~setup ~body ~teardown in
   "catchup" >:::[
     "common" >:: w test_common;
     "with_doubles" >:: w test_with_doubles;

src/node/collapser_test.ml

@@ -126,7 +126,7 @@ let teardown (dn, tlf_dir, head_dir) =
 
 
 let suite =
-  let wrapTest f = Extra.lwt_bracket setup f teardown
+  let wrapTest body = Lwt_extra.OUnit.bracket ~setup ~body ~teardown
   in
   "collapser_test" >:::[
     "collapse_until" >:: wrapTest test_collapse_until;

src/node/store_test.ml

@@ -18,7 +18,6 @@ limitations under the License.
 
 open OUnit
 open Lwt
-open Extra
 open Update
 open Simple_store
 open Store
@@ -136,7 +135,7 @@ let test_safe_insert_value_with_partial_value_update () =
 
 
 let suite =
-  let w f = lwt_bracket setup f teardown in
+  let w body = Lwt_extra.OUnit.bracket ~setup ~body ~teardown in
   "store" >:::[
     "safe_insert_value" >:: w test_safe_insert_value;
     "safe_insert_value_with_partial_value_update" >:: w test_safe_insert_value_with_partial_value_update;

src/paxos/multi_paxos_test.ml

@@ -21,7 +21,6 @@ open Lwt
 open MPMessage
 open Messaging
 open Multi_paxos
-open Extra
 open Update
 open Lwt_buffer
 open Master_type
@@ -507,7 +506,7 @@ let c2_fails = [ (fun (_msg,s,_t) -> s <> "c2")]
 
 
 open OUnit
-let w = lwt_test_wrap
+let w case () = Lwt_main.run (case ())
 let suite = "basic" >::: [
     "singleton_perfect" >:: w (test_generic build_perfect 1);
     "pair_perfect"  >:: w (test_generic build_perfect 2);

src/system/drop_master.ml

@@ -135,5 +135,5 @@ let make_suite base name w =
 
 
 let suite =
-  let w tn base f = Extra.lwt_bracket (setup tn Elected base) f teardown in
+  let w tn base body = Lwt_extra.OUnit.bracket ~setup:(setup tn Elected base) ~body ~teardown in
   make_suite 8000 "drop_master" w

src/system/single.ml

@@ -652,11 +652,11 @@ let make_suite base name w =
 
 let force_master =
   let make_master tn n = Forced (_node_name tn  n) in
-  let w tn base f = Extra.lwt_bracket (setup make_master tn base) f teardown in
+  let w tn base body = Lwt_extra.OUnit.bracket ~setup:(setup make_master tn base) ~body ~teardown in
   make_suite 4000 "force_master" w
 
 
 let elect_master =
   let make_master _tn _ = Elected in
-  let w tn base f = Extra.lwt_bracket (setup make_master tn base) f teardown in
+  let w tn base body = Lwt_extra.OUnit.bracket ~setup:(setup make_master tn base) ~body ~teardown in
   make_suite 6000 "elect_master" w

src/system/startup.ml

@@ -428,7 +428,7 @@ let interrupted_election () =
   Lwt_list.iter_s (_dump_tlc ~tlcs)   [node2; node3]
 
 
-let w f = Extra.lwt_bracket setup f teardown
+let w body = Lwt_extra.OUnit.bracket ~setup ~body ~teardown
 
 let suite = "startup" >:::[
     "post_failure" >:: w post_failure;

src/tlog/compression_test.ml

@@ -18,7 +18,6 @@ limitations under the License.
 
 open Compression
 open Lwt
-open Extra
 open OUnit
 open Tlogwriter
 open Update
@@ -58,7 +57,7 @@ let test_compress_file which () =
   OUnit.assert_equal md5 md5';
   Lwt.return()
 
-let w= lwt_test_wrap
+let w case () = Lwt_main.run (case ())
 
 let snappy =
   let archive_name x = x ^ ".tlx"

src/tlog/tlogcollection_test.ml

@@ -17,7 +17,6 @@ limitations under the License.
 
 
 open OUnit
-open Extra
 open Lwt
 open Update
 open Tlogcollection
@@ -248,7 +247,7 @@ let test_validate_corrupt_1 (dn, tlf_dir, factory) =
   >>= fun () ->
   Lwt.return ()
 
-let wrap factory test (name:string) = lwt_bracket (setup factory name) test teardown
+let wrap factory body (name:string) = Lwt_extra.OUnit.bracket ~setup:(setup factory name) ~body ~teardown
 
 let create_test_tlc dn = Mem_tlogcollection.make_mem_tlog_collection dn None None ~fsync:false ~fsync_tlog_dir:false
 

src/tools/lwt_extra.ml

@@ -71,3 +71,19 @@ end = struct
                     else ();
                 Lwt.return ()))
 end
+
+module OUnit : sig
+  val bracket :  setup:(unit -> 'a Lwt.t)
+              -> body:('a -> unit Lwt.t)
+              -> teardown:('a -> unit Lwt.t)
+              -> unit
+              -> unit
+end = struct
+  let bracket ~setup ~body ~teardown () =
+    Lwt_main.run begin
+      setup () >>= fun x ->
+      Lwt.map ignore (Lwt.finalize
+        (fun () -> body x)
+        (fun () -> teardown x))
+    end
+end

src/tools/lwt_socket_test.ml

@@ -59,11 +59,11 @@ let test_leak () =
       end
     ]
 
-let wrap t =
-  Extra.lwt_bracket
-    (fun () -> Lwt.return ())
-    t
-    (fun () -> Lwt.return ())
+let wrap body =
+  Lwt_extra.OUnit.bracket
+    ~setup:(fun () -> Lwt.return ())
+    ~body
+    ~teardown:(fun () -> Lwt.return ())
 let suite = "server_socket" >:::[
     "leak" >:: wrap test_leak
   ]