From 1842922fbecccfed639b82efa11982f498d71f3b Mon Sep 17 00:00:00 2001 From: Alex Tercete Date: Wed, 30 Apr 2025 15:45:08 +0100 Subject: [PATCH 1/2] feat(runner): handle testing env vars Previously, this was done in `labgrid_test`. Doing it in the runner is more flexible (Python insted of Starlark), and also paves the way for custom runner support in `labgrid_test`. --- bazel/labgrid/runner/runner.py | 30 ++++++++++++++++++++++++++++-- labgrid/test/rule.bzl | 14 +------------- 2 files changed, 29 insertions(+), 15 deletions(-) diff --git a/bazel/labgrid/runner/runner.py b/bazel/labgrid/runner/runner.py index 8905e2a..f9de04a 100644 --- a/bazel/labgrid/runner/runner.py +++ b/bazel/labgrid/runner/runner.py @@ -101,7 +101,7 @@ class Runner: Transfers will be attempted on a non zero exit code. """ - for download in downloads: + for download in downloads + self._default_downloads(): remote, local = self._resolve(download) # FIXME: Check if file exists with `@ape//ape:test` instead of assuming any error is due to a missing file try: @@ -112,7 +112,9 @@ class Runner: def run(self, cmd: str, env: Mapping[str, str] = {}) -> int: """Run a command on the device with given environment variables.""" - cmd = f"cd {self._exec_root} && {self._tools.env(env, cmd)}" + cmd = ( + f"cd {self._exec_root} && {self._tools.env(self._default_env() | env, cmd)}" + ) out, err, code = self._run(cmd, check=False) for line in out: stdout.write(f"{line}{linesep}") @@ -168,3 +170,27 @@ class Runner: def _read_path(self, cmd): out, _, _ = self._run(cmd) return PurePath(linesep.join(out).rstrip()) + + @staticmethod + def _default_env(): + inherit_env = ["BAZEL_TEST"] + + env = { + "XML_OUTPUT_FILE": "results.xml", + } + + return {k: environ.get(k) for k in inherit_env} | { + k: v for k, v in env.items() if k in environ + } + + @staticmethod + def _default_downloads(): + get = { + "XML_OUTPUT_FILE": "results.xml", + } + + return [ + FileTransfer(v, environ.get(k), optional=True) + for k, v in get.items() + if k in environ + ] diff --git a/labgrid/test/rule.bzl b/labgrid/test/rule.bzl index d8ad436..d88cb20 100644 --- a/labgrid/test/rule.bzl +++ b/labgrid/test/rule.bzl @@ -65,22 +65,11 @@ def _env(pair): key, value = pair return ["--env", "{}={}".format(key, value)] -def _get(pair): - key, value = pair - return ["--get", "{}:{}".format(key, value)] - def _files(targets): return [f for t in targets for f in t.files.to_list()] def implementation(ctx): - env = _run_env(ctx.attr.src) | {k: _expand(ctx, ctx.attr.data, v) for k, v in ctx.attr.env.items()} | { - "XML_OUTPUT_FILE": "results.xml", - "BAZEL_TEST": "1", - } - - get = { - "results.xml?": "${XML_OUTPUT_FILE}", - } + env = _run_env(ctx.attr.src) | {k: _expand(ctx, ctx.attr.data, v) for k, v in ctx.attr.env.items()} arguments = ctx.actions.declare_file("{}.args".format(ctx.label.name)) args = ctx.actions.args() @@ -88,7 +77,6 @@ def implementation(ctx): args.add(to_rlocation_path(ctx, ctx.executable._run)) args.add_all(["--program-prefix", paths.dirname(ctx.executable.src.short_path)]) args.add_all(env.items(), map_each = _env) - args.add_all(get.items(), map_each = _get) args.add("--") args.add(to_rlocation_path(ctx, ctx.executable.src)) ctx.actions.write(output = arguments, content = args) -- GitLab From d9002a2499ef91aba14f649ffaa94d43cf2df852 Mon Sep 17 00:00:00 2001 From: Alex Tercete Date: Wed, 30 Apr 2025 17:04:06 +0100 Subject: [PATCH 2/2] feat(runner): support test sharding --- bazel/labgrid/runner/runner.py | 6 ++++-- e2e/test/BUILD.bazel | 27 +++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/bazel/labgrid/runner/runner.py b/bazel/labgrid/runner/runner.py index f9de04a..d42c658 100644 --- a/bazel/labgrid/runner/runner.py +++ b/bazel/labgrid/runner/runner.py @@ -173,13 +173,14 @@ class Runner: @staticmethod def _default_env(): - inherit_env = ["BAZEL_TEST"] + inherit_env = ["BAZEL_TEST", "TEST_TOTAL_SHARDS", "TEST_SHARD_INDEX"] env = { "XML_OUTPUT_FILE": "results.xml", + "TEST_SHARD_STATUS_FILE": "test.shard", } - return {k: environ.get(k) for k in inherit_env} | { + return {k: environ.get(k) for k in inherit_env if k in environ} | { k: v for k, v in env.items() if k in environ } @@ -187,6 +188,7 @@ class Runner: def _default_downloads(): get = { "XML_OUTPUT_FILE": "results.xml", + "TEST_SHARD_STATUS_FILE": "test.shard", } return [ diff --git a/e2e/test/BUILD.bazel b/e2e/test/BUILD.bazel index 5aa1ead..4e243af 100644 --- a/e2e/test/BUILD.bazel +++ b/e2e/test/BUILD.bazel @@ -137,3 +137,30 @@ labgrid_test( }, platform = "@rules_labgrid//labgrid/platform:localhost", ) + +write_file( + name = "assert-sharding", + out = "sharding.py", + content = [ + "from os import environ", + "from pathlib import Path", + "assert int(environ['TEST_TOTAL_SHARDS']) == 2", + "assert int(environ['TEST_SHARD_INDEX']) in [0, 1]", + "Path(environ['TEST_SHARD_STATUS_FILE']).touch()", + ], +) + +py_binary( + name = "inner-sharding", + srcs = ["sharding.py"], + main = "sharding.py", + deps = ["@rules_python//python/runfiles"], +) + +labgrid_test( + name = "sharding", + size = "small", + src = ":inner-sharding", + platform = "@rules_labgrid//labgrid/platform:localhost", + shard_count = 2, +) -- GitLab