From 56ff8e14a832785f6c3b7601ddefaac619d6693d Mon Sep 17 00:00:00 2001
From: Zhongpeng Lin <zplin@uber.com>
Date: Thu, 22 Jun 2023 16:12:32 -0700
Subject: [PATCH] Allow passing deps in py_pytest_main

---
 examples/pytest/BUILD.bazel   | 6 ++++--
 examples/pytest/README.md     | 5 ++++-
 py/private/py_pytest_main.bzl | 4 +++-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/examples/pytest/BUILD.bazel b/examples/pytest/BUILD.bazel
index 93984361..5e7c0c36 100644
--- a/examples/pytest/BUILD.bazel
+++ b/examples/pytest/BUILD.bazel
@@ -1,7 +1,10 @@
 load("@rules_python//python:defs.bzl", "py_test")
 load("@//py:defs.bzl", "py_pytest_main")
 
-py_pytest_main(name = "__test__")
+py_pytest_main(
+    name = "__test__",
+    deps = ["@pypi_pytest//:pkg"],
+)
 
 py_test(
     name = "pytest_test",
@@ -13,6 +16,5 @@ py_test(
     main = ":__test__.py",
     deps = [
         ":__test__",
-        "@pypi_pytest//:pkg",
     ],
 )
diff --git a/examples/pytest/README.md b/examples/pytest/README.md
index 65050265..caebfd28 100644
--- a/examples/pytest/README.md
+++ b/examples/pytest/README.md
@@ -12,7 +12,10 @@ After importing this repository into your `WORKSPACE`, add to your Bazel package
 ```python
 load("@aspect_rules_py//py:defs.bzl", "py_pytest_main")
 
-py_pytest_main(name = "__test__")
+py_pytest_main(
+    name = "__test__",
+    deps = ["@pypi_pytest//:pkg"], # change this to the pytest target in your repo.
+)
 ```
 
 When using this repository together with the Gazelle extension for Python from rules_python, Gazelle
diff --git a/py/private/py_pytest_main.bzl b/py/private/py_pytest_main.bzl
index ecc133f9..6ae594b9 100644
--- a/py/private/py_pytest_main.bzl
+++ b/py/private/py_pytest_main.bzl
@@ -53,12 +53,13 @@ _py_pytest_main = rule(
     },
 )
 
-def py_pytest_main(name, py_library = default_py_library, **kwargs):
+def py_pytest_main(name, py_library = default_py_library, deps = [], **kwargs):
     """py_pytest_main wraps the template rendering target and the final py_library.
 
     Args:
         name: The name of the runable target that updates the test entry file.
         py_library: Use this attribute to override the default py_library rule.
+        deps: A list containing the pytest library target, e.g., @pypi_pytest//:pkg.
         **kwargs: The extra arguments passed to the template rendering target.
     """
 
@@ -79,4 +80,5 @@ def py_pytest_main(name, py_library = default_py_library, **kwargs):
         srcs = [test_main],
         tags = tags,
         visibility = visibility,
+        deps = deps,
     )