diff --git a/packages/bigframes/noxfile.py b/packages/bigframes/noxfile.py index f3b23f7e7d18..788471e0587d 100644 --- a/packages/bigframes/noxfile.py +++ b/packages/bigframes/noxfile.py @@ -28,6 +28,11 @@ import nox import nox.sessions +PROJECT_ID_OVERRIDE = os.getenv("BIGFRAMES_TEST_PROJECT") +ENV_OVERRIDES = ( + {"GOOGLE_CLOUD_PROJECT": PROJECT_ID_OVERRIDE} if PROJECT_ID_OVERRIDE else {} +) + RUFF_VERSION = "ruff==0.14.14" MYPY_VERSION = "mypy==1.15.0" @@ -367,11 +372,7 @@ def run_system( ) pytest_cmd.extend(extra_pytest_options) - session.run( - *pytest_cmd, - *session.posargs, - test_folder, - ) + session.run(*pytest_cmd, *session.posargs, test_folder, env=ENV_OVERRIDES) @nox.session(python="3.12") @@ -636,6 +637,7 @@ def prerelease(session: nox.sessions.Session, tests_path, extra_pytest_options=( tests_path, *extra_pytest_options, *session.posargs, + env=ENV_OVERRIDES, ) @@ -670,7 +672,7 @@ def system_prerelease(session: nox.sessions.Session): @nox.session(python=COLAB_AND_BQ_STUDIO_PYTHON_VERSIONS) def notebook(session: nox.Session): - google_cloud_project = os.getenv("GOOGLE_CLOUD_PROJECT") + google_cloud_project = PROJECT_ID_OVERRIDE or os.getenv("GOOGLE_CLOUD_PROJECT") if not google_cloud_project: session.error( "Set GOOGLE_CLOUD_PROJECT environment variable to run notebook session." @@ -766,6 +768,7 @@ def notebook(session: nox.Session): "python", CURRENT_DIRECTORY / "scripts" / "notebooks_fill_params.py", *notebooks, + env=ENV_OVERRIDES, ) processes = [] @@ -778,8 +781,7 @@ def notebook(session: nox.Session): ) if multi_process_mode: process = multiprocessing.Process( - target=session.run, - args=args, + target=session.run, args=args, kwargs={"env": ENV_OVERRIDES} ) process.start() processes.append(process) @@ -787,7 +789,7 @@ def notebook(session: nox.Session): # process to avoid potential race conditions。 time.sleep(1) else: - session.run(*args) + session.run(*args, env=ENV_OVERRIDES) for notebook, regions in notebooks_reg.items(): for region in regions: @@ -802,6 +804,7 @@ def notebook(session: nox.Session): process = multiprocessing.Process( target=session.run, args=region_args, + kwargs={"env": ENV_OVERRIDES}, ) process.start() processes.append(process) @@ -809,7 +812,7 @@ def notebook(session: nox.Session): # process to avoid potential race conditions。 time.sleep(1) else: - session.run(*region_args) + session.run(*region_args, env=ENV_OVERRIDES) for process in processes: process.join() @@ -826,6 +829,7 @@ def notebook(session: nox.Session): "scripts/run_and_publish_benchmark.py", "--notebook", "--publish-benchmarks=notebooks/", + env=ENV_OVERRIDES, ) @@ -889,6 +893,7 @@ def benchmark(session: nox.Session): "scripts/run_and_publish_benchmark.py", f"--benchmark-path={benchmark}", f"--iterations={args.iterations}", + env=ENV_OVERRIDES, ) finally: session.run( @@ -897,6 +902,7 @@ def benchmark(session: nox.Session): f"--publish-benchmarks={base_path}", f"--iterations={args.iterations}", f"--output-csv={args.output_csv}", + env=ENV_OVERRIDES, ) @@ -917,7 +923,7 @@ def release_dry_run(session): @nox.session(python=DEFAULT_PYTHON_VERSION) def cleanup(session): """Clean up stale and/or temporary resources in the test project.""" - google_cloud_project = os.getenv("GOOGLE_CLOUD_PROJECT") + google_cloud_project = PROJECT_ID_OVERRIDE or os.getenv("GOOGLE_CLOUD_PROJECT") cleanup_options = [] if google_cloud_project: cleanup_options.append(f"--project-id={google_cloud_project}")