Integration test reorg and build sharding (#540)

## Summary

This PR implements a custom pytest plugin for test sharding, allowing
tests to be automatically distributed across multiple CI jobs for
parallel execution. This replaces the previous manual directory-based
test splitting approach with a more flexible, hash-based distribution
system. Additionally, integration tests have been reorganized into
top-level folders that group tests by the client type and setup
requirements needed.

This should bring down the total CI runtime to ~8 minutes or less.

## Changes

### Core Implementation

- **New pytest plugin** (`tests/pytest_shard.py`):
- Implements `pytest_addoption` hook to add `--splits` and `--group`
command-line options
- Implements `pytest_collection_modifyitems` hook to filter tests based
on shard assignment
- Uses hash-based distribution (MD5 hash of test node ID) for
deterministic test assignment
- Supports environment variables `PYTEST_SPLITS` and `PYTEST_GROUP` as
alternatives to command-line options
  - Includes validation for shard parameters with helpful error messages

- **Plugin registration** (`tests/conftest.py`):
  - Registers the plugin globally so it's available for all test runs
  - Plugin is automatically loaded when running pytest

### CI Workflow Updates

- **Updated `.github/workflows/testing-integration.yaml`**:
- Replaced manual directory-based test splitting with automatic sharding
- `rest_sync` tests: Now uses 8 shards (was manually split by directory)
- `rest_asyncio` tests: Now uses 5 shards (was manually split by
directory)
- `grpc` tests: No sharding (runs all tests in single job, including
`tests/integration/rest_sync/db/data` with `USE_GRPC='true'`)

- **Updated `.github/actions/run-integration-test/action.yaml`**:
  - Added `pytest_splits` and `pytest_group` input parameters
  - Updated test execution to pass sharding arguments when provided

### Test Reorganization

- **Integration tests reorganized by client type**
(`tests/integration/`):
- **`rest_sync/`**: Tests using the synchronous REST client
(`Pinecone()`)
    - Uses standard `Index()` objects for database operations
- Supports optional GRPC mode via `USE_GRPC='true'` environment variable
- Contains subdirectories for `db/` (control and data operations),
`inference/`, and `admin/` tests
- **`rest_asyncio/`**: Tests using the asynchronous REST client
(`Pinecone().IndexAsyncio()`)
    - Uses async fixtures and `IndexAsyncio()` objects
    - Requires `pytest-asyncio` for async test execution
- Contains subdirectories for `db/` (control and data operations) and
`inference/` tests
  - **`grpc/`**: Tests using the GRPC client (`PineconeGRPC()`)
    - Uses `PineconeGRPC()` client and `GRPCIndex` objects
    - Contains `db/data/` tests for GRPC-specific functionality
- This organization makes it clear which client type each test requires
and simplifies fixture setup

### Bug Fixes

- **Fixed race condition in test cleanup**
(`tests/integration/rest_sync/db/control/pod/conftest.py`):
- Added `NotFoundException` handling in `attempt_delete_index` function
- Prevents teardown errors when index is deleted between `has_index`
check and `describe_index` call

### Testing

- **Unit tests** (`tests/unit/test_pytest_shard.py`):
  - Tests for hash-based distribution logic
  - Tests for validation and error handling
  - Tests for deterministic shard assignment
  - Tests for edge cases (single shard, environment variables, etc.)
- Tests gracefully handle `testdir` limitations (plugin loading in
isolated environments)

### Documentation

- **Updated `docs/maintainers/testing-guide.md`**:
  - Added "Test Sharding" section with usage examples
  - Documented command-line options and environment variables
  - Explained how sharding works and its use in CI
  - Documented actual shard counts used in CI workflows
  - Fixed broken link to `testing-integration.yaml`

## Benefits

1. **Automatic test distribution**: Tests are automatically distributed
across shards using a deterministic hash algorithm, eliminating the need
to manually maintain directory-based splits
2. **Better load balancing**: Hash-based distribution ensures more even
test distribution across shards compared to directory-based splitting
3. **Easier maintenance**: No need to manually update CI workflows when
test files are added, removed, or reorganized
4. **Flexibility**: Shard counts can be easily adjusted in CI workflows
without code changes
5. **Deterministic**: Same test always goes to the same shard, making
debugging easier
6. **Clear test organization**: Tests are grouped by client type, making
it immediately clear which setup and fixtures are needed for each test
7. **Simplified fixture management**: Each client type has its own
`conftest.py` with appropriate fixtures, reducing complexity and
potential conflicts

## Usage

### Command-line
```sh
pytest tests/integration/rest_sync --splits=8 --group=1
```

### Environment variables
```sh
export PYTEST_SPLITS=8
export PYTEST_GROUP=1
pytest tests/integration/rest_sync
```

## Testing

- Plugin works correctly in real pytest environment
- CI workflows updated and ready for use

## Notes

- The plugin is automatically available when running pytest (no
installation needed)
- Shard counts in CI can be adjusted based on test suite size and CI
capacity

Author: jhamon

.durations_grpc

@@ -17,11 +17,15 @@ inputs:
   dimension:
     description: 'The dimension of the index'
     required: false
-    default: '3'
+    default: ''
   metric:
     description: 'The metric of the index'
     required: false
     default: 'cosine'
+  vector_type:
+    description: 'The type of the index'
+    required: false
+    default: 'dense'
   PINECONE_API_KEY:
     description: 'The Pinecone API key'
     required: true
@@ -36,6 +40,10 @@ outputs:
     description: 'The name of the index, including randomized suffix'
     value: ${{ steps.create-index.outputs.index_name }}
 
+  index_host:
+    description: 'The host of the index'
+    value: ${{ steps.create-index.outputs.index_host }}
+
 runs:
   using: 'composite'
   steps:
@@ -52,5 +60,6 @@ runs:
         NAME_PREFIX: ${{ inputs.name_prefix }}
         REGION: ${{ inputs.region }}
         CLOUD: ${{ inputs.cloud }}
+        VECTOR_TYPE: ${{ inputs.vector_type }}
         DIMENSION: ${{ inputs.dimension }}
         METRIC: ${{ inputs.metric }}

.durations_rest_asyncio

@@ -1,9 +1,9 @@
 import os
-import re
 import random
 import string
-from datetime import datetime
+import uuid
 from pinecone import Pinecone
+from datetime import datetime
 
 
 def read_env_var(name):
@@ -22,39 +22,9 @@ def write_gh_output(name, value):
         print(f"{name}={value}", file=fh)
 
 
-def generate_index_name(test_name: str) -> str:
-    github_actor = os.getenv("GITHUB_ACTOR", None)
-    user = os.getenv("USER", None)
-    index_owner = github_actor or user
-
-    formatted_date = datetime.now().strftime("%Y%m%d-%H%M%S%f")[:-3]
-
-    github_job = os.getenv("GITHUB_JOB", None)
-
-    if test_name.startswith("test_"):
-        test_name = test_name[5:]
-
-    # Remove trailing underscore, if any
-    if test_name.endswith("_"):
-        test_name = test_name[:-1]
-
-    name_parts = [index_owner, formatted_date, github_job, test_name]
-    index_name = "-".join([x for x in name_parts if x is not None])
-
-    # Remove invalid characters
-    replace_with_hyphen = re.compile(r"[\[\(_,\s]")
-    index_name = re.sub(replace_with_hyphen, "-", index_name)
-    replace_with_empty = re.compile(r"[\]\)\.]")
-    index_name = re.sub(replace_with_empty, "", index_name)
-
-    max_length = 45
-    index_name = index_name[:max_length]
-
-    # Trim final character if it is not alphanumeric
-    if index_name.endswith("_") or index_name.endswith("-"):
-        index_name = index_name[:-1]
-
-    return index_name.lower()
+def generate_index_name(name_prefix: str) -> str:
+    name = name_prefix.lower() + "-" + str(uuid.uuid4())
+    return name[:45]
 
 
 def get_tags():
@@ -74,15 +44,35 @@ def get_tags():
 
 def main():
     pc = Pinecone(api_key=read_env_var("PINECONE_API_KEY"))
-    index_name = generate_index_name(read_env_var("NAME_PREFIX") + random_string(20))
+    index_name = generate_index_name(read_env_var("NAME_PREFIX"))
+    dimension_var = read_env_var("DIMENSION")
+    if dimension_var is not None and dimension_var != "":
+        dimension = int(dimension_var)
+    else:
+        dimension = None
+
+    vector_type_var = read_env_var("VECTOR_TYPE")
+    if vector_type_var is not None and vector_type_var != "":
+        vector_type = vector_type_var
+    else:
+        vector_type = None
+
+    metric = read_env_var("METRIC")
+    cloud = read_env_var("CLOUD")
+    region = read_env_var("REGION")
+    tags = get_tags()
+
     pc.create_index(
         name=index_name,
-        metric=read_env_var("METRIC"),
-        dimension=int(read_env_var("DIMENSION")),
-        spec={"serverless": {"cloud": read_env_var("CLOUD"), "region": read_env_var("REGION")}},
-        tags=get_tags(),
+        metric=metric,
+        dimension=dimension,
+        vector_type=vector_type,
+        tags=tags,
+        spec={"serverless": {"cloud": cloud, "region": region}},
     )
+    description = pc.describe_index(name=index_name)
     write_gh_output("index_name", index_name)
+    write_gh_output("index_host", description.host)
 
 
 if __name__ == "__main__":

.durations_rest_sync

@@ -14,11 +14,29 @@ inputs:
   PINECONE_ADDITIONAL_HEADERS:
     description: 'Additional headers to send with the request'
     required: false
-    default: '{"sdk-test-suite": "pinecone-python-client"}'
+    default: '{"sdk-test-suite": "pinecone-python-client", "x-environment": "preprod-aws-0"}'
   use_grpc:
     description: 'Whether to use gRPC or REST'
     required: false
     default: 'false'
+  PINECONE_CLIENT_ID:
+    description: 'The client ID to use for admin tests'
+    required: false
+  PINECONE_CLIENT_SECRET:
+    description: 'The client secret to use for admin tests'
+    required: false
+  INDEX_HOST_DENSE:
+    description: 'The host of the dense index for db data tests'
+    required: false
+  INDEX_HOST_SPARSE:
+    description: 'The host of the sparse index for db data tests'
+    required: false
+  pytest_splits:
+    description: 'Number of shards to split tests into (for test sharding)'
+    required: false
+  pytest_group:
+    description: 'Which shard to run (1-indexed, for test sharding)'
+    required: false
 
 runs:
   using: 'composite'
@@ -33,9 +51,23 @@ runs:
     - name: Run tests
       id: run-tests
       shell: bash
-      run: poetry run pytest tests/integration/${{ inputs.test_suite }} --retries 2 --retry-delay 35 -s -vv --log-cli-level=DEBUG --durations=20
+      run: |
+        PYTEST_ARGS=""
+        if [ -n "${{ inputs.pytest_splits }}" ] && [ -n "${{ inputs.pytest_group }}" ]; then
+          PYTEST_ARGS="--splits=${{ inputs.pytest_splits }} --group=${{ inputs.pytest_group }}"
+        fi
+        poetry run pytest ${{ inputs.test_suite }} \
+          $PYTEST_ARGS \
+          --retries 2 \
+          --retry-delay 35 \
+          --log-cli-level=DEBUG \
+          --durations=25 \
+          -s -vv
       env:
         PINECONE_API_KEY: ${{ steps.decrypt-api-key.outputs.decrypted_secret }}
         PINECONE_ADDITIONAL_HEADERS: ${{ inputs.PINECONE_ADDITIONAL_HEADERS }}
+        PINECONE_CLIENT_ID: ${{ inputs.PINECONE_CLIENT_ID }}
+        PINECONE_CLIENT_SECRET: ${{ inputs.PINECONE_CLIENT_SECRET }}
         USE_GRPC: ${{ inputs.use_grpc }}
-        SKIP_WEIRD: 'true'
+        INDEX_HOST_DENSE: ${{ inputs.INDEX_HOST_DENSE }}
+        INDEX_HOST_SPARSE: ${{ inputs.INDEX_HOST_SPARSE }}

.github/actions/index-create/action.yml

@@ -21,6 +21,10 @@ inputs:
     description: 'Python version to use'
     required: true
     default: '3.10'
+  enable_cache:
+    description: 'Enable caching of Poetry dependencies and virtual environment'
+    required: true
+    default: 'true'
 
 runs:
   using: 'composite'
@@ -33,6 +37,21 @@ runs:
     - name: Install Poetry
       uses: snok/install-poetry@v1
 
+    - name: Get Poetry cache directory
+      if: ${{ inputs.enable_cache == 'true' }}
+      id: poetry-cache
+      shell: bash
+      run: |
+        echo "dir=$(poetry config cache-dir)" >> $GITHUB_OUTPUT
+
+    - name: Cache Poetry dependencies
+      if: ${{ inputs.enable_cache == 'true' }}
+      uses: actions/cache@v4
+      id: restore-cache-poetry
+      with:
+        path: ${{ steps.poetry-cache.outputs.dir }}
+        key: poetry-${{ runner.os }}-${{ inputs.python_version }}-${{ hashFiles('poetry.lock') }}-grpc-${{ inputs.include_grpc }}-asyncio-${{ inputs.include_asyncio }}-dev-${{ inputs.include_dev }}-types-${{ inputs.include_types }}
+
     - name: Install dependencies
       shell: bash
       env:

.github/actions/index-create/create.py

@@ -30,6 +30,7 @@ runs:
         include_types: false
         include_asyncio: true
         python_version: ${{ inputs.python_version }}
+        enable_cache: 'false'
 
     - name: 'Install aiohttp ${{ inputs.aiohttp_version }}'
       run: 'poetry add aiohttp==${{ inputs.aiohttp_version }}'

.github/actions/run-integration-test/action.yaml

@@ -38,6 +38,7 @@ runs:
         include_grpc: true
         include_types: false
         python_version: ${{ inputs.python_version }}
+        enable_cache: 'false'
 
     - name: Install grpcio ${{ inputs.grpcio_version }}
       run: poetry add grpcio==${{ inputs.grpcio_version }}

.github/actions/setup-poetry/action.yml

@@ -29,6 +29,7 @@ runs:
         include_grpc: false
         include_types: false
         python_version: ${{ inputs.python_version }}
+        enable_cache: 'false'
 
     - name: 'Install urllib3 ${{ matrix.urllib3-version }}'
       run: 'poetry add urllib3==${{ matrix.urllib3-version }}'