diff --git a/bazel/labgrid/runner/runner.py b/bazel/labgrid/runner/runner.py index 8905e2af6836fdebcb9b4eab147be0bdea1775fc..d42c658b1dd2ddb91519c435f99125f8b60d0c95 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,29 @@ 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", "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 if k in environ} | { + k: v for k, v in env.items() if k in environ + } + + @staticmethod + def _default_downloads(): + get = { + "XML_OUTPUT_FILE": "results.xml", + "TEST_SHARD_STATUS_FILE": "test.shard", + } + + return [ + FileTransfer(v, environ.get(k), optional=True) + for k, v in get.items() + if k in environ + ] diff --git a/e2e/test/BUILD.bazel b/e2e/test/BUILD.bazel index 5aa1ead5601bdea1a96dd3ff4c25cbf278aff19c..4e243aff1611e66adc5d31f7e7162bfca41b1008 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, +) diff --git a/labgrid/test/rule.bzl b/labgrid/test/rule.bzl index d8ad436f7e138058a378ea1e61f6e7f981442ac1..d88cb2098f15afb8ef65244760c76bc7788c353a 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)