Source file expect_test_collector.ml

open Expect_test_common
module List = ListLabels

module Test_outcome = struct
  type t =
    { file_digest : File.Digest.t
    ; location : File.Location.t
    ; expectations : Expectation.Raw.t list
    ; uncaught_exn_expectation : Expectation.Raw.t option
    ; saved_output : (File.Location.t * string) list
    ; trailing_output : string
    ; upon_unreleasable_issue : Expect_test_config_types.Upon_unreleasable_issue.t
    ; uncaught_exn : (exn * Printexc.raw_backtrace) option
    }
end

let tests_run : Test_outcome.t list ref = ref []

let protect ~finally ~f =
  match f () with
  | x ->
    finally ();
    x
  | exception e ->
    finally ();
    raise e
;;

module Current_file = struct
  let current = ref None

  let set ~absolute_filename =
    match !current with
    | None -> current := Some absolute_filename
    | Some _ -> failwith "Expect_test_collector.set: already set"
  ;;

  let unset () =
    match !current with
    | Some _ -> current := None
    | None -> failwith "Expect_test_collector.unset: not set"
  ;;

  let get () =
    match !current with
    | Some fn -> fn
    | None -> failwith "Expect_test_collector.get: not set"
  ;;
end

module Instance = struct
  type t =
    { mutable saved : (File.Location.t * int) list
    ; chan : out_channel
    ; filename : File.Name.t
    }

  external before_test
    :  output:out_channel
    -> stdout:out_channel
    -> stderr:out_channel
    -> unit
    = "expect_test_collector_before_test"

  external after_test
    :  stdout:out_channel
    -> stderr:out_channel
    -> unit
    = "expect_test_collector_after_test"

  external pos_out : out_channel -> int = "caml_out_channel_pos_fd"

  let get_position () = pos_out stdout

  let create () =
    let filename = Filename.temp_file "expect-test" "output" in
    let chan = open_out_bin filename in
    before_test ~output:chan ~stdout ~stderr;
    { chan; filename = File.Name.of_string filename; saved = [] }
  ;;

  let relative_filename t = File.Name.relative_to ~dir:(File.initial_dir ()) t.filename

  let with_ic fname ~f =
    let ic = open_in_bin fname in
    protect ~finally:(fun () -> close_in ic) ~f:(fun () -> f ic)
  ;;

  let current_test : (File.Location.t * t) option ref = ref None

  let get_current () =
    match !current_test with
    | Some (_, t) -> t
    | None -> failwith "Expect_test_collector.Instance.get_current called outside a test."
  ;;

  let save_output_without_flush t location =
    let pos = get_position () in
    t.saved <- (location, pos) :: t.saved
  ;;

  let save_and_return_output_without_flush t location =
    let pos = get_position () in
    let prev_pos =
      match t.saved with
      | [] -> 0
      | (_, prev_pos) :: _ -> prev_pos
    in
    t.saved <- (location, pos) :: t.saved;
    flush t.chan;
    let len = pos - prev_pos in
    with_ic (relative_filename t) ~f:(fun ic ->
      seek_in ic prev_pos;
      really_input_string ic len)
  ;;
end

let basic_flush () =
  Format.pp_print_flush Format.std_formatter ();
  Format.pp_print_flush Format.err_formatter ();
  Stdlib.flush Stdlib.stdout;
  Stdlib.flush Stdlib.stderr
;;

let save_and_return_output location =
  let instance = Instance.get_current () in
  basic_flush ();
  Instance.save_and_return_output_without_flush instance location
;;

module Make (C : Expect_test_config_types.S) = struct
  let ( >>= ) t f = C.IO_flush.bind t ~f
  let return = C.IO_flush.return

  module C = struct
    include C

    let flush () =
      basic_flush ();
      C.IO_flush.return ()
    ;;
  end

  module Instance_io : sig
    val save_output : File.Location.t -> unit C.IO_flush.t
    val save_and_return_output : File.Location.t -> string C.IO_flush.t

    val exec
      :  file_digest:File.Digest.t
      -> location:File.Location.t
      -> expectations:Expectation.Raw.t list
      -> uncaught_exn_expectation:Expectation.Raw.t option
      -> f:(unit -> unit C.IO_run.t)
      -> unit
  end = struct
    open Instance

    let extract_output_and_sanitize ic len =
      let s = really_input_string ic len |> C.sanitize in
      if not (Check_backtraces.contains_backtraces s)
      then s
      else
        Expect_test_config_types.Upon_unreleasable_issue
        .message_when_expectation_contains_backtrace
          C.upon_unreleasable_issue
        ^ s
    ;;

    let get_outputs_and_cleanup t =
      let last_ofs = get_position () in
      after_test ~stdout ~stderr;
      close_out t.chan;
      let fname = relative_filename t in
      protect
        ~finally:(fun () -> Sys.remove fname)
        ~f:(fun () ->
          with_ic fname ~f:(fun ic ->
            let ofs, outputs =
              List.fold_left
                (List.rev t.saved)
                ~init:(0, [])
                ~f:(fun (ofs, acc) (loc, next_ofs) ->
                  let s = extract_output_and_sanitize ic (next_ofs - ofs) in
                  next_ofs, (loc, s) :: acc)
            in
            let trailing_output = extract_output_and_sanitize ic (last_ofs - ofs) in
            List.rev outputs, trailing_output))
    ;;

    let save_output location =
      let t = get_current () in
      C.flush ()
      >>= fun () ->
      save_output_without_flush t location;
      return ()
    ;;

    let save_and_return_output location =
      let t = get_current () in
      C.flush () >>= fun () -> return (save_and_return_output_without_flush t location)
    ;;

    let () =
      Stdlib.at_exit (fun () ->
        match !current_test with
        | None -> ()
        | Some (loc, t) ->
          let blocks, trailing = get_outputs_and_cleanup t in
          Printf.eprintf
            "File %S, line %d, characters %d-%d:\n\
             Error: program exited while expect test was running!\n\
             Output captured so far:\n\
             %!"
            (File.Name.to_string loc.filename)
            loc.line_number
            (loc.start_pos - loc.line_start)
            (loc.end_pos - loc.line_start);
          List.iter blocks ~f:(fun (_, s) -> Printf.eprintf "%s%!" s);
          Printf.eprintf "%s%!" trailing)
    ;;

    let rec final_flush ?(count = 0) k =
      let max_attempts = 10 in
      C.flush ()
      >>= fun () ->
      if C.flushed ()
      then k ~append:""
      else if count = max_attempts
      then
        k
          ~append:
            (Printf.sprintf
               "\n\
                STOPPED COLLECTING OUTPUT AFTER %d FLUSHING ATTEMPS\n\
                THERE MUST BE A BACKGROUND JOB PRINTING TO STDOUT\n"
               max_attempts)
      else final_flush ~count:(count + 1) k
    ;;

    let exec ~file_digest ~location ~expectations ~uncaught_exn_expectation ~f =
      let t = create () in
      current_test := Some (location, t);
      let finally uncaught_exn =
        C.run (fun () ->
          C.IO_flush.to_run
            (final_flush (fun ~append ->
               current_test := None;
               let saved_output, trailing_output = get_outputs_and_cleanup t in
               tests_run
               := { file_digest
                  ; location
                  ; expectations
                  ; uncaught_exn_expectation
                  ; saved_output
                  ; trailing_output = trailing_output ^ append
                  ; upon_unreleasable_issue = C.upon_unreleasable_issue
                  ; uncaught_exn
                  }
                  :: !tests_run;
               return ())))
      in
      match C.run f with
      | () -> finally None
      | exception exn ->
        let bt = Printexc.get_raw_backtrace () in
        finally (Some (exn, bt))
    ;;
  end

  let save_output = Instance_io.save_output
  let save_and_return_output = Instance_io.save_and_return_output

  let run
        ~file_digest
        ~(location : File.Location.t)
        ~absolute_filename:defined_in
        ~description
        ~tags
        ~expectations
        ~uncaught_exn_expectation
        ~inline_test_config
        f
    =
    Ppx_inline_test_lib.Runtime.test
      ~config:inline_test_config
      ~descr:
        (lazy
          (match description with
           | None -> ""
           | Some s -> s))
      ~tags
      ~filename:(File.Name.to_string location.filename)
      ~line_number:location.line_number
      ~start_pos:(location.start_pos - location.line_start)
      ~end_pos:(location.end_pos - location.line_start)
      (fun () ->
         let registering_tests_for = Current_file.get () in
         if defined_in <> registering_tests_for
         then
           Printf.ksprintf
             failwith
             "Trying to run an expect test from the wrong file.\n\
              - test declared at %s:%d\n\
              - trying to run it from %s\n"
             defined_in
             location.line_number
             registering_tests_for
         else (
           (* To avoid capturing not-yet flushed data of the stdout buffer *)
           C.run (fun () -> C.IO_flush.to_run (C.flush ()));
           Instance_io.exec
             ~file_digest
             ~location
             ~expectations
             ~uncaught_exn_expectation
             ~f;
           true))
  ;;
end
[@@inline never]

let tests_run () =
  (* We prepend tests when we encounter them, so reverse the list to reinstate order *)
  List.rev !tests_run
;;