Collect test types in CI in parallel

.github/workflows/ci.yml

@@ -843,6 +843,11 @@ jobs:
       RUNS_ON: "${{needs.build-info.outputs.runs-on}}"
       BACKEND: sqlite
       PYTHON_MAJOR_MINOR_VERSION: "${{needs.build-info.outputs.default-python-version}}"
+      TEST_TYPES: "${{needs.build-info.outputs.test-types}}"
+      FULL_TESTS_NEEDED: "${{needs.build-info.outputs.full-tests-needed}}"
+      DEBUG_RESOURCES: "${{needs.build-info.outputs.debug-resources}}"
+      JOB_ID: "test-pytest-collection"
+      COVERAGE: "false"
     steps:
       - name: Cleanup repo
         run: docker run -v "${GITHUB_WORKSPACE}:/workspace" -u 0:0 bash -c "rm -rf /workspace/*"
@@ -858,7 +863,7 @@ jobs:
         env:
           DEBUG_RESOURCES: ${{needs.build-info.outputs.debug-resources}}
       - name: "Tests Pytest collection"
-        run: breeze shell "python /opt/airflow/scripts/in_container/test_pytest_collection.py"
+        run: breeze testing tests --run-in-parallel --collect-only
       - name: "Fix ownership"
         run: breeze ci fix-ownership
         if: always()
@@ -904,7 +909,7 @@ jobs:
       - name: "Tests: ${{matrix.python-version}}:${{needs.build-info.outputs.test-types}}"
         run: breeze testing tests --run-in-parallel
       - name: "Tests ARM Pytest collection: ${{matrix.python-version}}"
-        run: breeze shell "python /opt/airflow/scripts/in_container/test_pytest_collection.py" arm
+        run: breeze testing tests --run-in-parallel --collect-only --remove-arm-packages
       - name: "Post Tests: ${{matrix.python-version}}:${{needs.build-info.outputs.test-types}}"
         uses: ./.github/actions/post_tests
 
@@ -992,7 +997,7 @@ jobs:
         run: breeze testing tests --run-in-parallel
       - name: "Tests ARM Pytest collection: ${{matrix.python-version}}"
         if: env.MYSQL_VERSION != '5.7'
-        run: breeze shell "python /opt/airflow/scripts/in_container/test_pytest_collection.py" arm
+        run: breeze testing tests --run-in-parallel --collect-only --remove-arm-packages
       - name: "Post Tests: ${{matrix.python-version}}:${{needs.build-info.outputs.test-types}}"
         uses: ./.github/actions/post_tests
 
@@ -1076,7 +1081,7 @@ jobs:
       - name: "Tests: ${{matrix.python-version}}:${{needs.build-info.outputs.test-types}}"
         run: breeze testing tests --run-in-parallel
       - name: "Tests ARM Pytest collection: ${{matrix.python-version}}"
-        run: breeze shell "python /opt/airflow/scripts/in_container/test_pytest_collection.py" arm
+        run: breeze testing tests --run-in-parallel --collect-only --remove-arm-packages
       - name: "Post Tests: ${{matrix.python-version}}:${{needs.build-info.outputs.test-types}}"
         uses: ./.github/actions/post_tests
 

Dockerfile.ci

@@ -912,6 +912,21 @@ if [[ ${ENABLE_TEST_COVERAGE:="false"} == "true" ]]; then
     )
 fi
 
+if [[ ${COLLECT_ONLY:="false"} == "true" ]]; then
+    EXTRA_PYTEST_ARGS+=(
+        "--collect-only"
+        "-qqqq"
+        "--disable-warnings"
+    )
+fi
+
+if [[ ${REMOVE_ARM_PACKAGES:="false"} == "true" ]]; then
+    # Test what happens if we do not have ARM packages installed.
+    # This is useful to see if pytest collection works without ARM packages which is important
+    # for the MacOS M1 users running tests in their ARM machines with `breeze testing tests` command
+    python "${IN_CONTAINER_DIR}/remove_arm_packages.py"
+fi
+
 declare -a SELECTED_TESTS CLI_TESTS API_TESTS PROVIDERS_TESTS CORE_TESTS WWW_TESTS \
     ALL_TESTS ALL_PRESELECTED_TESTS ALL_OTHER_TESTS
 

dev/breeze/src/airflow_breeze/commands/testing_commands.py

@@ -131,6 +131,8 @@ def _run_test(
     if db_reset:
         env_variables["DB_RESET"] = "true"
     env_variables["TEST_TYPE"] = exec_shell_params.test_type
+    env_variables["COLLECT_ONLY"] = str(exec_shell_params.collect_only).lower()
+    env_variables["REMOVE_ARM_PACKAGES"] = str(exec_shell_params.remove_arm_packages).lower()
     env_variables["SKIP_PROVIDER_TESTS"] = str(exec_shell_params.skip_provider_tests).lower()
     if "[" in exec_shell_params.test_type and not exec_shell_params.test_type.startswith("Providers"):
         get_console(output=output).print(
@@ -367,6 +369,18 @@ def run_tests_in_parallel(
     is_flag=True,
     envvar="UPGRADE_BOTO",
 )
+@click.option(
+    "--collect-only",
+    help="Collect tests only, do not run them.",
+    is_flag=True,
+    envvar="COLLECT_ONLY",
+)
+@click.option(
+    "--remove-arm-packages",
+    help="Removes arm packages from the image to test if ARM collection works",
+    is_flag=True,
+    envvar="REMOVE_ARM_PACKAGES",
+)
 @option_verbose
 @option_dry_run
 @click.argument("extra_pytest_args", nargs=-1, type=click.UNPROCESSED)
@@ -391,6 +405,8 @@ def command_for_tests(
     mount_sources: str,
     extra_pytest_args: tuple,
     upgrade_boto: bool,
+    collect_only: bool,
+    remove_arm_packages: bool,
 ):
     docker_filesystem = get_filesystem_type("/var/lib/docker")
     get_console().print(f"Docker filesystem: {docker_filesystem}")
@@ -406,6 +422,8 @@ def command_for_tests(
         forward_ports=False,
         test_type=test_type,
         upgrade_boto=upgrade_boto,
+        collect_only=collect_only,
+        remove_arm_packages=remove_arm_packages,
     )
     rebuild_or_pull_ci_image_if_needed(command_params=exec_shell_params)
     cleanup_python_generated_files()

dev/breeze/src/airflow_breeze/commands/testing_commands_config.py

@@ -27,6 +27,7 @@
             "options": [
                 "--test-type",
                 "--test-timeout",
+                "--collect-only",
                 "--db-reset",
                 "--backend",
                 "--python",
@@ -54,6 +55,7 @@
                 "--image-tag",
                 "--mount-sources",
                 "--upgrade-boto",
+                "--remove-arm-packages",
             ],
         },
     ],

dev/breeze/src/airflow_breeze/params/shell_params.py

@@ -76,6 +76,7 @@ class ShellParams:
     backend: str = ALLOWED_BACKENDS[0]
     base_branch: str = "main"
     ci: bool = False
+    collect_only: bool = False
     db_reset: bool = False
     dev_mode: bool = False
     extra_args: tuple = ()
@@ -102,6 +103,7 @@ class ShellParams:
     platform: str = DOCKER_DEFAULT_PLATFORM
     postgres_version: str = ALLOWED_POSTGRES_VERSIONS[0]
     python: str = ALLOWED_PYTHON_MAJOR_MINOR_VERSIONS[0]
+    remove_arm_packages: bool = False
     skip_environment_initialization: bool = False
     skip_constraints: bool = False
     start_airflow: str = "false"

dev/breeze/src/airflow_breeze/utils/docker_command_utils.py

@@ -596,6 +596,7 @@ def update_expected_environment_variables(env: dict[str, str]) -> None:
     set_value_to_default_if_not_set(env, "CI_TARGET_BRANCH", AIRFLOW_BRANCH)
     set_value_to_default_if_not_set(env, "CI_TARGET_REPO", APACHE_AIRFLOW_GITHUB_REPOSITORY)
     set_value_to_default_if_not_set(env, "COMMIT_SHA", commit_sha())
+    set_value_to_default_if_not_set(env, "COLLECT_ONLY", "false")
     set_value_to_default_if_not_set(env, "DB_RESET", "false")
     set_value_to_default_if_not_set(env, "DEFAULT_BRANCH", AIRFLOW_BRANCH)
     set_value_to_default_if_not_set(env, "ENABLED_SYSTEMS", "")
@@ -610,6 +611,7 @@ def update_expected_environment_variables(env: dict[str, str]) -> None:
     set_value_to_default_if_not_set(env, "LOAD_EXAMPLES", "false")
     set_value_to_default_if_not_set(env, "PACKAGE_FORMAT", ALLOWED_PACKAGE_FORMATS[0])
     set_value_to_default_if_not_set(env, "PYTHONDONTWRITEBYTECODE", "true")
+    set_value_to_default_if_not_set(env, "REMOVE_ARM_PACKAGES", "false")
     set_value_to_default_if_not_set(env, "RUN_SYSTEM_TESTS", "false")
     set_value_to_default_if_not_set(env, "RUN_TESTS", "false")
     set_value_to_default_if_not_set(env, "SKIP_ENVIRONMENT_INITIALIZATION", "false")

images/breeze/output-commands-hash.txt

@@ -60,5 +60,5 @@ stop:e5aa686b4e53707ced4039d8414d5cd6
 testing:docker-compose-tests:b86c044b24138af0659a05ed6331576c
 testing:helm-tests:94a442e7f3f63b34c4831a84d165690a
 testing:integration-tests:225ddb6243cce5fc64f4824b87adfd98
-testing:tests:6abb9e771e9af02d352ef1cafccf260e
-testing:dec193f57ad7ac385b785dc68adf1f66
+testing:tests:953ce65d99acee99014061b2fab7b8a3
+testing:bc492ec01c4de953f28a66d0e8f50bc9