diff --git a/tests/unit/ws/t_ws.ml b/tests/unit/ws/t_ws.ml
index 48d3b09b..61f474d9 100644
--- a/tests/unit/ws/t_ws.ml
+++ b/tests/unit/ws/t_ws.ml
@@ -6,9 +6,9 @@ let read_file file : string =
 
 let apply_masking = Tiny_httpd_ws.Private_.apply_masking
 
-let decode ~key b =
+let decode ?(mask_offset = 0) ~key b =
   let buf = Bytes.copy b in
-  apply_masking ~mask_key:key buf 0 (Bytes.length buf);
+  apply_masking ~mask_key:key ~mask_offset buf 0 (Bytes.length buf);
   buf
 
 let () =
diff --git a/tests/unit/ws/t_ws_q.ml b/tests/unit/ws/t_ws_q.ml
index 6243a973..1e17cb50 100644
--- a/tests/unit/ws/t_ws_q.ml
+++ b/tests/unit/ws/t_ws_q.ml
@@ -2,9 +2,9 @@ open Test_util
 
 let apply_masking = Tiny_httpd_ws.Private_.apply_masking
 
-let decode ~key b =
+let decode ?(mask_offset = 0) ~key b =
   let buf = Bytes.copy b in
-  apply_masking ~mask_key:key buf 0 (Bytes.length buf);
+  apply_masking ~mask_key:key ~mask_offset buf 0 (Bytes.length buf);
   buf
 
 let ppb b = Printf.sprintf "%S" (Bytes.unsafe_to_string b)
@@ -13,12 +13,15 @@ let () =
   add_qcheck
   @@ QCheck.Test.make ~count:10_000
        Q.(
-         pair (bytes_of_size (Gen.return 4)) (bytes_of_size Gen.(0 -- 6000))
+         triple
+           (bytes_of_size (Gen.return 4))
+           (option small_nat)
+           (bytes_of_size Gen.(0 -- 6000))
          (* |> Q.add_stat ("b.size", fun (_k, b) -> Bytes.length b) *)
-         |> Q.add_shrink_invariant (fun (k, _) -> Bytes.length k = 4))
-       (fun (key, b) ->
+         |> Q.add_shrink_invariant (fun (k, _, _) -> Bytes.length k = 4))
+       (fun (key, mask_offset, b) ->
          Q.assume (Bytes.length key = 4);
-         let b2 = decode ~key b in
+         let b2 = decode ~key ?mask_offset b in
 
          let is_zero =
            Bytes.equal key (Bytes.unsafe_of_string "\x00\x00\x00\x00")
@@ -32,7 +35,7 @@ let () =
                (ppb key) (ppb b) (ppb b2)
          );
 
-         let b3 = decode ~key b2 in
+         let b3 = decode ~key ?mask_offset b2 in
          assert (Bytes.equal b b3);
          true)