Implement ref with cleanup

src/ReactDOM.re

@@ -487,9 +487,13 @@ module Ref = {
   type t = domRef;
   type currentDomRef = React.ref(Js.nullable(Dom.element));
   type callbackDomRef = Js.nullable(Dom.element) => unit;
+  type callbackRefWithCleanup = (Js.nullable(Dom.element), unit) => unit;
 
   external domRef: currentDomRef => domRef = "%identity";
   external callbackDomRef: callbackDomRef => domRef = "%identity";
+  external callbackRefWithCleanup:
+    [@mel.unwrap] (callbackRefWithCleanup => domRef) =
+    "%identity";
 };
 
 /* This list isn't exhaustive. We'll add more as we go. */

src/ReactDOM.rei

@@ -488,9 +488,12 @@ module Ref: {
   type t = domRef;
   type currentDomRef = React.ref(Js.nullable(Dom.element));
   type callbackDomRef = Js.nullable(Dom.element) => unit;
+  type callbackRefWithCleanup = (Js.nullable(Dom.element), unit) => unit;
 
   external domRef: currentDomRef => domRef = "%identity";
   external callbackDomRef: callbackDomRef => domRef = "%identity";
+  external callbackRefWithCleanup: callbackRefWithCleanup => domRef =
+    "%identity";
 };
 
 /* This list isn't exhaustive. We'll add more as we go. */

test/Ref__test.re

@@ -31,5 +31,54 @@ describe("ref", () => {
       | None => failwith("No element found")
       };
     expect(content)->toBe("Click me");
-  })
+  });
+
+  test("with callback ref", () => {
+    let callbackRef = Mock.fn();
+
+    let refWithCleanup =
+      ReactDOM.Ref.callbackDomRef(_ref => {ignore(callbackRef())});
+
+    let _container =
+      ReactTestingLibrary.render(
+        <Button_with_ref ref=refWithCleanup>
+          {React.string("Click me")}
+        </Button_with_ref>,
+      );
+
+    expect(callbackRef->Mock.getMock)->toHaveBeenCalled();
+  });
+
+  let (let.await) = (p, f) => Js.Promise.then_(f, p);
+
+  testPromise("with callback ref and cleanup", finish => {
+    let callbackRef = Mock.fnWithImplementation(_ => ());
+    let callbackCleanupRef = Mock.fnWithImplementation(_ => ());
+
+    let refWithCleanup =
+      ReactDOM.Ref.callbackRefWithCleanup(_ref => {
+        callbackRef();
+        () => {
+          callbackCleanupRef();
+        };
+      });
+
+    let.await _ =
+      ReactTestingLibrary.actAsync(() => {
+        let _container =
+          ReactTestingLibrary.render(
+            <Button_with_ref ref=refWithCleanup>
+              {React.string("Click me")}
+            </Button_with_ref>,
+          );
+        Js.Promise.resolve();
+      });
+
+    expect(callbackRef->Mock.getMock)->toHaveBeenCalledTimes(1);
+    expect(callbackCleanupRef->Mock.getMock)->toHaveBeenCalledTimes(0);
+
+    ReactTestingLibrary.cleanup();
+    expect(callbackCleanupRef->Mock.getMock)->toHaveBeenCalledTimes(1);
+    finish();
+  });
 });

test/jest/Expect.re

@@ -20,6 +20,12 @@ external toBeLessThanOrEqual: (t('a), 'a) => unit = "toBeLessThanOrEqual";
 [@mel.send]
 external toHaveLength: (t(array('a)), 'a) => unit = "toHaveLength";
 
+[@mel.send]
+external toHaveBeenCalled: (t('a), unit) => unit = "toHaveBeenCalled";
+[@mel.send]
+external toHaveBeenCalledTimes: (t('a), int) => unit =
+  "toHaveBeenCalledTimes";
+
 [@mel.get]
 external rejects: t(Js.Promise.t('a)) => t(unit => 'a) = "rejects";
 

test/jest/Mock.re

@@ -5,7 +5,7 @@ let undefined: undefined = Js.Undefined.empty;
 type mock('a);
 
 type t('a) = mock('a);
-[@mel.scope "jest"] external fn: unit => 'a = "fn";
+[@mel.scope "jest"] external fn: [@mel.unwrap] (unit => 'a) = "fn";
 [@mel.scope "jest"] external fnWithImplementation: 'a => 'a = "fn";
 [@mel.scope "jest"] external mockModule: string => unit = "mock";
 external getMock: 'a => t('a) = "%identity";