2828import nox
2929import nox .sessions
3030
31+ PROJECT_ID_OVERRIDE = os .getenv ("BIGFRAMES_TEST_PROJECT" )
32+ ENV_OVERRIDES = (
33+ {"GOOGLE_CLOUD_PROJECT" : PROJECT_ID_OVERRIDE } if PROJECT_ID_OVERRIDE else {}
34+ )
35+
3136RUFF_VERSION = "ruff==0.14.14"
3237MYPY_VERSION = "mypy==1.15.0"
3338
@@ -367,11 +372,7 @@ def run_system(
367372 )
368373
369374 pytest_cmd .extend (extra_pytest_options )
370- session .run (
371- * pytest_cmd ,
372- * session .posargs ,
373- test_folder ,
374- )
375+ session .run (* pytest_cmd , * session .posargs , test_folder , env = ENV_OVERRIDES )
375376
376377
377378@nox .session (python = "3.12" )
@@ -636,6 +637,7 @@ def prerelease(session: nox.sessions.Session, tests_path, extra_pytest_options=(
636637 tests_path ,
637638 * extra_pytest_options ,
638639 * session .posargs ,
640+ env = ENV_OVERRIDES ,
639641 )
640642
641643
@@ -670,7 +672,7 @@ def system_prerelease(session: nox.sessions.Session):
670672
671673@nox .session (python = COLAB_AND_BQ_STUDIO_PYTHON_VERSIONS )
672674def notebook (session : nox .Session ):
673- google_cloud_project = os .getenv ("GOOGLE_CLOUD_PROJECT" )
675+ google_cloud_project = PROJECT_ID_OVERRIDE or os .getenv ("GOOGLE_CLOUD_PROJECT" )
674676 if not google_cloud_project :
675677 session .error (
676678 "Set GOOGLE_CLOUD_PROJECT environment variable to run notebook session."
@@ -766,6 +768,7 @@ def notebook(session: nox.Session):
766768 "python" ,
767769 CURRENT_DIRECTORY / "scripts" / "notebooks_fill_params.py" ,
768770 * notebooks ,
771+ env = ENV_OVERRIDES ,
769772 )
770773
771774 processes = []
@@ -778,16 +781,15 @@ def notebook(session: nox.Session):
778781 )
779782 if multi_process_mode :
780783 process = multiprocessing .Process (
781- target = session .run ,
782- args = args ,
784+ target = session .run , args = args , kwargs = {"env" : ENV_OVERRIDES }
783785 )
784786 process .start ()
785787 processes .append (process )
786788 # Adding a small delay between starting each
787789 # process to avoid potential race conditions。
788790 time .sleep (1 )
789791 else :
790- session .run (* args )
792+ session .run (* args , env = ENV_OVERRIDES )
791793
792794 for notebook , regions in notebooks_reg .items ():
793795 for region in regions :
@@ -802,14 +804,15 @@ def notebook(session: nox.Session):
802804 process = multiprocessing .Process (
803805 target = session .run ,
804806 args = region_args ,
807+ kwargs = {"env" : ENV_OVERRIDES },
805808 )
806809 process .start ()
807810 processes .append (process )
808811 # Adding a small delay between starting each
809812 # process to avoid potential race conditions。
810813 time .sleep (1 )
811814 else :
812- session .run (* region_args )
815+ session .run (* region_args , env = ENV_OVERRIDES )
813816
814817 for process in processes :
815818 process .join ()
@@ -826,6 +829,7 @@ def notebook(session: nox.Session):
826829 "scripts/run_and_publish_benchmark.py" ,
827830 "--notebook" ,
828831 "--publish-benchmarks=notebooks/" ,
832+ env = ENV_OVERRIDES ,
829833 )
830834
831835
@@ -889,6 +893,7 @@ def benchmark(session: nox.Session):
889893 "scripts/run_and_publish_benchmark.py" ,
890894 f"--benchmark-path={ benchmark } " ,
891895 f"--iterations={ args .iterations } " ,
896+ env = ENV_OVERRIDES ,
892897 )
893898 finally :
894899 session .run (
@@ -897,6 +902,7 @@ def benchmark(session: nox.Session):
897902 f"--publish-benchmarks={ base_path } " ,
898903 f"--iterations={ args .iterations } " ,
899904 f"--output-csv={ args .output_csv } " ,
905+ env = ENV_OVERRIDES ,
900906 )
901907
902908
@@ -917,7 +923,7 @@ def release_dry_run(session):
917923@nox .session (python = DEFAULT_PYTHON_VERSION )
918924def cleanup (session ):
919925 """Clean up stale and/or temporary resources in the test project."""
920- google_cloud_project = os .getenv ("GOOGLE_CLOUD_PROJECT" )
926+ google_cloud_project = PROJECT_ID_OVERRIDE or os .getenv ("GOOGLE_CLOUD_PROJECT" )
921927 cleanup_options = []
922928 if google_cloud_project :
923929 cleanup_options .append (f"--project-id={ google_cloud_project } " )
0 commit comments