chore: add test sharding to unit tests

.github/workflows/unittest.yml

@@ -13,11 +13,40 @@ permissions:
   contents: read
 
 jobs:
+  initialize:
+    runs-on: ubuntu-latest
+    outputs:
+      matrix: ${{ steps.set-matrix.outputs.matrix }}
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+    - 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 }}
+      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@v6
@@ -34,22 +63,41 @@ jobs:
       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@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:

ci/get_package_shards.py

@@ -0,0 +1,94 @@
+import os
+import subprocess
+import json
+import math
+import sys
+
+def get_packages():
+    subdirs = ['packages']
+    packages = []
+    for subdir in subdirs:
+        if not os.path.exists(subdir):
+            continue
+        # Use the same sorting as the shell script
+        pkg_dirs = [os.path.join(subdir, d) + '/' for d in os.listdir(subdir) if os.path.isdir(os.path.join(subdir, d))]
+        packages.extend(sorted(pkg_dirs))
+    return packages
+
+def get_packages_to_test():
+    build_type = os.environ.get('BUILD_TYPE', 'presubmit')
+    target_branch = os.environ.get('TARGET_BRANCH', 'main')
+
+    all_packages = get_packages()
+
+    if build_type == 'presubmit':
+        git_diff_arg = f"origin/{target_branch}..."
+    elif build_type == 'continuous':
+        git_diff_arg = "HEAD~.."
+    else:
+        return all_packages
+
+    try:
+        res = subprocess.check_output(['git', 'diff', '--name-only', git_diff_arg]).decode('utf-8')
+        changed_files = res.splitlines()
+    except subprocess.CalledProcessError:
+        return all_packages
+
+    to_test = []
+    for pkg in all_packages:
+        # Check if any changed file starts with the package path
+        if any(f.startswith(pkg) for f in changed_files):
+            to_test.append(pkg)
+
+    return to_test
+
+def group_packages(packages):
+    if not packages:
+        return []
+
+    num_packages = len(packages)
+
+    # 1. Only shard if > 10 packages are being touched
+    # 2. Only add a new shard if any shard would have > 10 packages.
+    # To guarantee that no shard contains more than 10 packages (when distributed evenly),
+    # we need S >= N / 10, which means S = ceil(N / 10).
+    num_shards = math.ceil(num_packages / 10)
+
+    # Ensure at least 1 shard if we have packages
+    num_shards = max(1, num_shards)
+
+    # 3. Top out at 10 shards
+    num_shards = min(10, num_shards)
+
+    # Distribute packages between them as evenly as possible
+    shard_size = math.ceil(num_packages / num_shards)
+
+    shards = []
+    for i in range(num_shards):
+        start = i * shard_size
+        end = min((i + 1) * shard_size, num_packages)
+        if start >= num_packages:
+            break
+        shard_packages = packages[start:end]
+        index = i + 1
+        name = f"Shard {index}"
+        num_in_shard = len(shard_packages)
+        # Calculate a descriptive range for step visibility
+        if len(shard_packages) == 1:
+            desc = shard_packages[0].strip('/').split('/')[-1]
+        else:
+            desc = f"{shard_packages[0].strip('/').split('/')[-1]}...{shard_packages[-1].strip('/').split('/')[-1]} ({num_in_shard} packages)"
+
+        shards.append({
+            "name": name,
+            "index": index,
+            "description": desc,
+            "packages": " ".join(shard_packages),
+            "is_sharded": num_shards > 1
+        })
+    return shards
+
+if __name__ == "__main__":
+    packages = get_packages_to_test()
+    shards = group_packages(packages)
+    print(json.dumps(shards))

ci/run_conditional_tests.sh

@@ -47,45 +47,73 @@ git config --global url."${PROJECT_ROOT}".insteadOf "https://github.com/googleap
 # A script file for running the test in a sub project.
 test_script="${PROJECT_ROOT}/ci/run_single_test.sh"
 
+if [ -n "${PACKAGE_LIST}" ]; then
+    echo "Using provided PACKAGE_LIST"
+    to_test=(${PACKAGE_LIST})
+    RETVAL=0
+    for d in ${to_test[@]}; do
+        echo "running test in ${d}"
+        pushd ${d}
+        set +e
+
+        # Ensure unique coverage file per package to avoid DataError
+        # when combining statement and branch coverage.
+        # Strip trailing slash from directory name for the filename.
+        pkg_name_clean=$(echo ${d} | sed 's|/$||' | sed 's|/|_|g')
+        export COVERAGE_FILE="${PROJECT_ROOT}/.coverage.${PY_VERSION}.${pkg_name_clean}"
+
+        ${test_script}
+        ret=$?
+        set -e
+        if [ ${ret} -ne 0 ]; then
+            RETVAL=${ret}
+        fi
+        popd
+    done
+    exit ${RETVAL}
+fi
+
 if [[ ${BUILD_TYPE} == "presubmit" ]]; then
     # For presubmit build, we want to know the difference from the
     # common commit in the target branch.
     GIT_DIFF_ARG="origin/$TARGET_BRANCH..."
 
     # Then fetch enough history for finding the common commit.
-    git fetch origin "$TARGET_BRANCH" --deepen=200
+    git fetch origin "$TARGET_BRANCH" --deepen=200 || true
 
 elif [[ ${BUILD_TYPE} == "continuous" ]]; then
     # For continuous build, we want to know the difference in the last
     # commit. This assumes we use squash commit when merging PRs.
     GIT_DIFF_ARG="HEAD~.."
 
     # Then fetch one last commit for getting the diff.
-    git fetch origin "$TARGET_BRANCH" --deepen=1
+    git fetch origin "$TARGET_BRANCH" --deepen=1 || true
 
 else
     # Run everything.
     GIT_DIFF_ARG=""
 fi
 
-# Then detect changes in the test scripts.
-
-set +e
-git diff --quiet ${GIT_DIFF_ARG} ci
-changed=$?
-set -e
-
-# Now we have a fixed list, but we can change it to autodetect if
-# necessary.
-
+# Sharding logic (fallback for manual runs)
 subdirs=(
     packages
 )
+TOTAL_SHARDS="${TOTAL_SHARDS:-1}"
+SHARD_INDEX="${SHARD_INDEX:-1}"
+count=0
 
 RETVAL=0
 
 for subdir in ${subdirs[@]}; do
-    for d in `ls -d ${subdir}/*/`; do
+    # Sort the directories to ensure consistent sharding across jobs
+    for d in `ls -d ${subdir}/*/ | sort`; do
+        # Sharding logic: only process directories that belong to this shard
+        if (( count % TOTAL_SHARDS != SHARD_INDEX - 1 )); then
+            ((++count))
+            continue
+        fi
+        ((++count))
+
         should_test=false
         if [ -n "${GIT_DIFF_ARG}" ]; then
             echo "checking changes with 'git diff --quiet ${GIT_DIFF_ARG} ${d}'"
@@ -108,6 +136,13 @@ for subdir in ${subdirs[@]}; do
             pushd ${d}
             # Temporarily allow failure.
             set +e
+
+            # Ensure unique coverage file per package to avoid DataError
+            # when combining statement and branch coverage.
+            # Strip trailing slash from directory name for the filename.
+            pkg_name_clean=$(echo ${d} | sed 's|/$||' | sed 's|/|_|g')
+            export COVERAGE_FILE="${PROJECT_ROOT}/.coverage.${PY_VERSION}.${pkg_name_clean}"
+
             ${test_script}
             ret=$?
             set -e

packages/google-backstory/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-apiregistry/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-bigtable/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-ces/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-chronicle/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-dataproc/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-firestore/.coveragerc

@@ -2,6 +2,7 @@
 branch = True
 
 [report]
+fail_under = 99
 show_missing = True
 omit =
     google/cloud/firestore/__init__.py

packages/google-cloud-firestore/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-gke-multicloud/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-managedkafka/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-optimization/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/google-cloud-retail/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/sqlalchemy-spanner/.coveragerc

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+[run]
+branch = True
+omit =
+  google/__init__.py
+  google/cloud/sqlalchemy_spanner/requirements.py
+
+[report]
+fail_under = 0
+show_missing = True
+exclude_lines =
+    # Re-enable the standard pragma
+    pragma: NO COVER
+    # Ignore debug-only repr
+    def __repr__
+    # Ignore abstract methods
+    raise NotImplementedError
+omit =
+  */site-packages/*.py
+  google/__init__.py
+  google/cloud/sqlalchemy_spanner/requirements.py

packages/sqlalchemy-spanner/SHARD_TEST.txt

@@ -0,0 +1 @@
+Test change for sharding

packages/sqlalchemy-spanner/noxfile.py

@@ -85,6 +85,7 @@ class = StreamHandler
 UNIT_TEST_STANDARD_DEPENDENCIES = [
     "mock",
     "pytest",
+    "pytest-cov",
 ]
 
 UNIT_TEST_EXTERNAL_DEPENDENCIES = [
@@ -368,7 +369,18 @@ def unit(session, test_type):
             *UNIT_TEST_DEPENDENCIES,
         )
         session.install(".")
-        session.run("py.test", "--quiet", os.path.join("tests/unit"), *session.posargs)
+        session.run(
+            "py.test",
+            "--quiet",
+            "--cov=google.cloud.sqlalchemy_spanner",
+            "--cov=tests/unit",
+            "--cov-append",
+            "--cov-config=.coveragerc",
+            "--cov-report=",
+            "--cov-fail-under=0",
+            os.path.join("tests/unit"),
+            *session.posargs,
+        )
         return