From e3c094aaccd87f91c3ed34452ea684d717386451 Mon Sep 17 00:00:00 2001
From: Alex Eagle <alex@aspect.dev>
Date: Tue, 28 Jan 2025 13:01:27 -0800
Subject: [PATCH] feat: py_test macro generates a py_pytest_main

Provides much-needed syntax sugar for this common case.
---
 e2e/use_release/src/BUILD.bazel | 24 ++++++------------------
 examples/pytest/BUILD.bazel     | 24 +-----------------------
 py/defs.bzl                     | 17 +++++++++++++++--
 3 files changed, 22 insertions(+), 43 deletions(-)

diff --git a/e2e/use_release/src/BUILD.bazel b/e2e/use_release/src/BUILD.bazel
index 85fbd578..1614141e 100644
--- a/e2e/use_release/src/BUILD.bazel
+++ b/e2e/use_release/src/BUILD.bazel
@@ -1,4 +1,4 @@
-load("@aspect_rules_py//py:defs.bzl", "py_binary", "py_pytest_main", "py_test")
+load("@aspect_rules_py//py:defs.bzl", "py_binary", "py_test")
 
 py_binary(
     name = "main",
@@ -9,24 +9,12 @@ py_binary(
     main = "__main__.py",
 )
 
-# TODO(alex): land https://github.com/aspect-build/rules_py/pull/401 and shorten this
-py_pytest_main(
-    name = "__test__",
-    deps = [
-        "@pip//coverage",
-        "@pip//pytest",
-    ],
-)
-
 py_test(
     name = "test",
-    srcs = [
-        "my_test.py",
-        ":__test__",
-    ],
-    main = ":__test__.py",
-    deps = [
-        ":__test__",
-        ":main",
+    srcs = ["my_test.py"],
+    pytest_main_deps = [
+        "@pip//coverage",
+        "@pip//pytest",
     ],
+    deps = [":main"],
 )
diff --git a/examples/pytest/BUILD.bazel b/examples/pytest/BUILD.bazel
index d423cacd..90515c94 100644
--- a/examples/pytest/BUILD.bazel
+++ b/examples/pytest/BUILD.bazel
@@ -38,29 +38,7 @@ py_test(
 )
 
 py_test(
-    name = "nested/pytest",
-    srcs = [
-        "foo_test.py",
-        ":__test__",
-    ],
-    data = glob([
-        "fixtures/*.json",
-    ]),
-    env_inherit = ["FOO"],
-    imports = ["../.."],
-    main = ":__test__.py",
-    package_collisions = "warning",
-    deps = [
-        ":__test__",
-        ":lib",
-        "@pypi_ftfy//:pkg",
-        "@pypi_neptune//:pkg",
-        "@pypi_pytest//:pkg",
-    ],
-)
-
-py_test(
-    name = "sharding_test",
+    name = "sharded/test",
     srcs = [
         "__test__.py",
         "sharding_test.py",
diff --git a/py/defs.bzl b/py/defs.bzl
index dc0e1772..58c0b219 100644
--- a/py/defs.bzl
+++ b/py/defs.bzl
@@ -117,7 +117,7 @@ def py_binary(name, srcs = [], main = None, **kwargs):
 
     _py_binary_or_test(name = name, rule = _py_binary, srcs = srcs, main = main, resolutions = resolutions, **kwargs)
 
-def py_test(name, srcs = [], main = None, **kwargs):
+def py_test(name, srcs = [], main = None, pytest_main_deps = None, **kwargs):
     """Identical to [py_binary](./py_binary.md), but produces a target that can be used with `bazel test`.
 
     Args:
@@ -127,6 +127,8 @@ def py_test(name, srcs = [], main = None, **kwargs):
             Like rules_python, this is treated as a suffix of a file that should appear among the srcs.
             If absent, then `[name].py` is tried. As a final fallback, if the srcs has a single file,
             that is used as the main.
+        pytest_main_deps: List of labels. If set, generate a [py_pytest_main](#py_pytest_main) script and use it as the py_test entry point.
+            The supplied labels define the deps attribute for the generated py_pytest_main.
         **kwargs: additional named parameters to `py_binary_rule`.
     """
 
@@ -139,4 +141,15 @@ def py_test(name, srcs = [], main = None, **kwargs):
     if resolutions:
         resolutions = resolutions.to_label_keyed_dict()
 
-    _py_binary_or_test(name = name, rule = _py_test, srcs = srcs, main = main, resolutions = resolutions, **kwargs)
+    deps = kwargs.pop("deps", [])
+    if pytest_main_deps:
+        pytest_main_target = name + ".pytest_main"
+        main = pytest_main_target + ".py"
+        py_pytest_main(
+            name = pytest_main_target,
+            deps = pytest_main_deps,
+        )
+        srcs.append(main)
+        deps.append(pytest_main_target)
+
+    _py_binary_or_test(name = name, rule = _py_test, srcs = srcs, deps = deps, main = main, resolutions = resolutions, **kwargs)