chore: add test sharding to unit tests

.github/workflows/unittest.yml

@@ -13,12 +13,68 @@ permissions:
   contents: read
 
 jobs:
+  initialize:
+    runs-on: ubuntu-latest
+    outputs:
+      matrix: ${{ steps.set-matrix.outputs.matrix }}
+      is_full_run: ${{ steps.check-label.outputs.is_full_run }}
+    env:
+      MAX_SHARDS: 20
+      # Define weights for long-running unit tests to balance shard execution time
+      # Each weight is roughly 1 minute of expected execution time
+      # Default for unset packages is 1
+      PACKAGE_WEIGHTS: |
+        bigframes: 6
+        google-ai-generativelanguage: 4
+        google-auth: 5
+        google-cloud-compute: 12
+        google-cloud-compute-v1beta: 12
+        google-cloud-dialogflow: 6
+        google-cloud-dialogflow-cx: 6
+        google-cloud-discoveryengine: 8
+        google-cloud-retail: 5
+        google-shopping-merchant-accounts: 4
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+    - name: Check for unit_test:all_packages label
+      id: check-label
+      run: |
+        if [[ "${{ contains(github.event.pull_request.labels.*.name, 'unit_test:all_packages') }}" == "true" ]]; then
+          echo "is_full_run=true" >> $GITHUB_OUTPUT
+        else
+          echo "is_full_run=false" >> $GITHUB_OUTPUT
+        fi
+    - name: Setup Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.10"
+    - name: Get package shards
+      id: set-matrix
+      env:
+        BUILD_TYPE: presubmit
+        TARGET_BRANCH: ${{ github.base_ref || github.event.merge_group.base_ref }}
+        TEST_ALL_PACKAGES: ${{ steps.check-label.outputs.is_full_run }}
+        MAX_SHARDS: ${{ env.MAX_SHARDS }}
+        PACKAGE_WEIGHTS: ${{ env.PACKAGE_WEIGHTS }}
+      run: |
+        if [ -n "$TARGET_BRANCH" ]; then
+          git fetch origin "$TARGET_BRANCH" --deepen=200 || true
+        fi
+        echo "matrix=$(python3 ci/get_package_shards.py)" >> $GITHUB_OUTPUT
+
   unit:
+    needs: initialize
+    if: needs.initialize.outputs.matrix != '[]' && needs.initialize.outputs.matrix != ''
     runs-on: ubuntu-22.04
     strategy:
       fail-fast: true
       matrix:
         python: ['3.9', '3.10', "3.11", "3.12", "3.13", "3.14"]
+        package_shard: ${{ fromJson(needs.initialize.outputs.matrix) }}
+    name: ${{ matrix.package_shard.is_sharded && format('unit ({0}, {1})', matrix.python, matrix.package_shard.name) || format('unit ({0})', matrix.python) }}
     steps:
     - name: Checkout
       uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6
@@ -32,30 +88,51 @@ jobs:
       uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6
       with:
         python-version: ${{ matrix.python }}
+        cache: 'pip'
     - name: Install nox
       run: |
         python -m pip install --upgrade setuptools pip wheel
         python -m pip install nox
-    - name: Run unit tests
+    - name: Run unit tests for ${{ matrix.package_shard.description }}
       env:
-        COVERAGE_FILE: ${{ github.workspace }}/.coverage-${{ matrix.python }}
         BUILD_TYPE: presubmit
         TARGET_BRANCH: ${{ github.base_ref || github.event.merge_group.base_ref }}
         TEST_TYPE: unit
         PY_VERSION: ${{ matrix.python }}
+        PACKAGE_LIST: ${{ matrix.package_shard.packages }}
       run: |
         ci/run_conditional_tests.sh
     - name: Upload coverage results
       uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4
       with:
-        name: coverage-artifact-${{ matrix.python }}
-        path: .coverage-${{ matrix.python }}
+        name: coverage-artifact-${{ matrix.python }}-${{ matrix.package_shard.index }}
+        path: .coverage.${{ matrix.python }}.*
         include-hidden-files: true
 
+  all-tests:
+    needs: [initialize, unit]
+    if: always()
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check unit test results
+        run: |
+          # 1. Check initialize job
+          if [[ "${{ needs.initialize.result }}" != "success" ]]; then
+            echo "Error: The initialize job status was: ${{ needs.initialize.result }}"
+            exit 1
+          fi
+          # 2. Check unit test shards
+          if [[ "${{ needs.unit.result }}" != "success" && "${{ needs.unit.result }}" != "skipped" ]]; then
+            echo "Unit tests failed"
+            exit 1
+          fi
+          echo "All unit tests passed or were skipped"
+
   cover:
     runs-on: ubuntu-latest
     needs:
         - unit
+        - initialize
     steps:
     - name: Checkout
       uses: actions/checkout@34e114876b0b11c390a56381ad16ebd13914f8d5 # v4
@@ -71,10 +148,16 @@ jobs:
         python-version: "3.10"
     - name: Set number of files changes in packages directory
       id: packages
+      env:
+        TEST_ALL_PACKAGES: ${{ needs.initialize.outputs.is_full_run }}
       run: |
-        git diff HEAD~1 -- packages > /dev/null
-        num_files_changed=$(git diff HEAD~1 -- packages | wc -l | tr -d ' ')
-        echo "num_files_changed=${num_files_changed}" >> "$GITHUB_OUTPUT"
+        if [[ "${TEST_ALL_PACKAGES}" == "true" ]]; then
+          echo "num_files_changed=1" >> "$GITHUB_OUTPUT"
+        else
+          git diff HEAD~1 -- packages > /dev/null
+          num_files_changed=$(git diff HEAD~1 -- packages | wc -l | tr -d ' ')
+          echo "num_files_changed=${num_files_changed}" >> "$GITHUB_OUTPUT"
+        fi
     - name: Install coverage
       if: ${{ steps.packages.outputs.num_files_changed > 0 }}
       run: |
@@ -84,86 +167,14 @@ jobs:
       if: ${{ steps.packages.outputs.num_files_changed > 0 }}
       uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5
       with:
-        path: .coverage-results/
+        path: /dev/shm/.coverage-results/
+        merge-multiple: true
     - name: Report coverage results
       if: ${{ steps.packages.outputs.num_files_changed > 0 }}
       env:
         # TODO: default to 100% coverage after next gapic-generator release
         # https://github.com/googleapis/google-cloud-python/issues/17459
         DEFAULT_FAIL_UNDER: 99
+        TEST_ALL_PACKAGES: ${{ needs.initialize.outputs.is_full_run }}
       run: |
-        if [ -d .coverage-results ]; then
-          # Unzip any zipped coverage results
-          find .coverage-results -type f -name '*.zip' -exec unzip -o {} \;
-
-          # Find all coverage files and combine them.
-          # We find files starting with .coverage (excluding .coveragerc files and templates)
-          coverage_files=$(find .coverage-results . -type f -name '.coverage*' ! -name '.coveragerc*')
-          if [ -n "${coverage_files}" ]; then
-            coverage combine ${coverage_files}
-          else
-            echo "Error: No coverage files found to combine."
-            exit 1
-          fi
-
-          # Find all modified packages
-          modified_packages=$(git diff --name-only HEAD~1 -- packages | cut -d/ -f1,2 | sort -u)
-
-          failed_packages=()
-          passed_packages=()
-
-          for pkg in ${modified_packages}; do
-            if [ -d "${pkg}" ]; then
-              echo "============================================================"
-              echo "Evaluating coverage for package: ${pkg}"
-              echo "============================================================"
-
-              set +e
-              pushd "${pkg}" > /dev/null
-              if [ -f ".coveragerc" ]; then
-                echo "Using package-specific configuration: ${pkg}/.coveragerc"
-                # If fail_under is specified in the package-specific .coveragerc, coverage report
-                # will automatically enforce it. Otherwise, we enforce the default.
-                if grep -q "fail_under" ".coveragerc"; then
-                  COVERAGE_FILE=../../.coverage coverage report --include="$PWD/**"
-                else
-                  echo "No fail_under specified in ${pkg}/.coveragerc, enforcing default"
-                  COVERAGE_FILE=../../.coverage coverage report --include="$PWD/**" --fail-under="${DEFAULT_FAIL_UNDER}"
-                fi
-              else
-                echo "No .coveragerc found for ${pkg}, enforcing default"
-                COVERAGE_FILE=../../.coverage coverage report --include="$PWD/**" --fail-under="${DEFAULT_FAIL_UNDER}"
-              fi
-              status=$?
-              popd > /dev/null
-              set -e
-
-              if [ ${status} -ne 0 ]; then
-                failed_packages+=("${pkg}")
-              else
-                passed_packages+=("${pkg}")
-              fi
-            fi
-          done
-
-          echo "============================================================"
-          echo "Coverage Evaluation Summary"
-          echo "============================================================"
-          if [ ${#passed_packages[@]} -gt 0 ]; then
-            echo "Passed packages:"
-            for pkg in "${passed_packages[@]}"; do
-              echo "  - ${pkg}"
-            done
-          fi
-          if [ ${#failed_packages[@]} -gt 0 ]; then
-            echo "Failed packages:"
-            for pkg in "${failed_packages[@]}"; do
-              echo "  - ${pkg}"
-            done
-            exit 1
-          fi
-        else
-          echo "Error: No coverage results were downloaded from the unit test jobs."
-          echo "This usually means the unit tests did not run or failed to upload their coverage files."
-          exit 1
-        fi
+        ci/report_coverage.sh