diff --git a/utils/bazel/llvm-project-overlay/llvm/lit_test.bzl b/utils/bazel/llvm-project-overlay/llvm/lit_test.bzl
index 0af12a635855df..64528f28492f57 100644
--- a/utils/bazel/llvm-project-overlay/llvm/lit_test.bzl
+++ b/utils/bazel/llvm-project-overlay/llvm/lit_test.bzl
@@ -4,7 +4,7 @@
 """Rules for running lit tests."""
 
 load("@bazel_skylib//lib:paths.bzl", "paths")
-load("@rules_python//python:defs.bzl", "py_test")
+load("@rules_python//python:defs.bzl", _py_test = "py_test")
 
 def lit_test(
         name,
@@ -12,6 +12,7 @@ def lit_test(
         args = None,
         data = None,
         deps = None,
+        py_test = _py_test,
         **kwargs):
     """Runs a single test file with LLVM's lit tool.
 
@@ -22,6 +23,8 @@ def lit_test(
         Note that `-v` and the 'srcs' paths are added automatically.
       data: label list. Additional data dependencies of the test.
         Note that 'srcs' targets are added automatically.
+      deps: label list. List of targets the test depends on.
+      py_test: function. The py_test rule to use for the underlying test.
       **kwargs: additional keyword arguments.
 
     See https://llvm.org/docs/CommandGuide/lit.html for details on lit.