Make pytest rerun count configurable in CI workflows and address bandit errors (#4001)

strickvl · web-flow · commit 1bdfff46e227 · 2025-09-26T08:47:58.000+02:00
* Update security checks and configuration settings

This commit modifies the `pyproject.toml` to include Google-style docstring conventions and updates the Bandit security tool configuration to exclude specific directories. Additionally, it enhances the `check-security.sh` script to set environment variables for debugging and analytics, ensuring a more robust security check process.

Changes:
- Add Google-style docstring convention in `pyproject.toml`
- Update Bandit configuration to exclude migrations directory
- Modify `check-security.sh` to set environment variables for improved security checks

* Make pytest rerun count configurable in CI workflows

Allows manual workflow triggers to specify custom pytest rerun count while maintaining default of 3 for automated CI runs.

Changes:
- Add PYTEST_RERUNS environment variable support to test-coverage-xml.sh with validation
- Add optional 'reruns' input to all test workflows (unit and integration)
- Pass reruns input as PYTEST_RERUNS environment variable to test scripts
- Maintain backward compatibility: unset variables default to 3 reruns

This helps developers iterate faster on flaky tests by reducing reruns during manual debugging sessions while keeping the safety net of 3 retries for automated CI.

* Fix variable references in warning messages

Correct the warning messages to properly show the environment variable name
(PYTEST_RERUNS/PYTEST_RERUNS_DELAY) while displaying the actual validated
value that was found to be invalid.

* Update Bandit configuration and migration script

This commit updates the `pyproject.toml` to refine the Bandit security tool configuration by excluding the migrations directory from checks. Additionally, it modifies the migration script `502b4fa5fa88_adding_in_progress_to_runs.py` to ensure secure database operations by using the `nosec` comment to suppress specific security warnings.

These changes enhance the security posture of the project while maintaining focus on relevant code areas.

Changes:
- Exclude migrations directory in Bandit configuration
- Add `nosec` comment in migration script for secure execution

* Fix placement of the #nosec
diff --git a/.github/workflows/integration-test-fast-services.yml b/.github/workflows/integration-test-fast-services.yml
@@ -25,6 +25,11 @@ on:
         type: number
         required: false
         default: 30
+      reruns:
+        description: Pytest rerun count (0 disables)
+        type: number
+        required: false
+        default: 3
   workflow_dispatch:
     inputs:
       os:
@@ -72,6 +77,11 @@ on:
         type: number
         required: false
         default: 30
+      reruns:
+        description: Pytest rerun count (0 disables)
+        type: number
+        required: false
+        default: 3
 jobs:
   integration-tests-fast:
     name: integration-tests-fast
@@ -97,6 +107,7 @@ jobs:
       GCP_US_EAST4_SERVER_URL: ${{ secrets.GCP_US_EAST4_SERVER_URL }}
       GCP_US_EAST4_SERVER_USERNAME: ${{ secrets.GCP_US_EAST4_SERVER_USERNAME }}
       GCP_US_EAST4_SERVER_PASSWORD: ${{ secrets.GCP_US_EAST4_SERVER_PASSWORD }}
+      PYTEST_RERUNS: ${{ inputs.reruns }}
     if: ${{ ! startsWith(github.event.head_commit.message, 'GitBook:') }}
     defaults:
       run:
diff --git a/.github/workflows/integration-test-fast.yml b/.github/workflows/integration-test-fast.yml
@@ -25,6 +25,11 @@ on:
         type: number
         required: false
         default: 30
+      reruns:
+        description: Pytest rerun count (0 disables)
+        type: number
+        required: false
+        default: 3
   workflow_dispatch:
     inputs:
       os:
@@ -72,6 +77,11 @@ on:
         type: number
         required: false
         default: 30
+      reruns:
+        description: Pytest rerun count (0 disables)
+        type: number
+        required: false
+        default: 3
 jobs:
   integration-tests-fast:
     name: integration-tests-fast
@@ -97,6 +107,7 @@ jobs:
       GCP_US_EAST4_SERVER_URL: ${{ secrets.GCP_US_EAST4_SERVER_URL }}
       GCP_US_EAST4_SERVER_USERNAME: ${{ secrets.GCP_US_EAST4_SERVER_USERNAME }}
       GCP_US_EAST4_SERVER_PASSWORD: ${{ secrets.GCP_US_EAST4_SERVER_PASSWORD }}
+      PYTEST_RERUNS: ${{ inputs.reruns }}
     if: ${{ ! startsWith(github.event.head_commit.message, 'GitBook:') }}
     defaults:
       run:
diff --git a/.github/workflows/integration-test-slow-services.yml b/.github/workflows/integration-test-slow-services.yml
@@ -25,6 +25,11 @@ on:
         type: number
         required: false
         default: 30
+      reruns:
+        description: Pytest rerun count (0 disables)
+        type: number
+        required: false
+        default: 3
   workflow_dispatch:
     inputs:
       os:
@@ -71,6 +76,11 @@ on:
         type: number
         required: false
         default: 30
+      reruns:
+        description: Pytest rerun count (0 disables)
+        type: number
+        required: false
+        default: 3
 jobs:
   integration-tests-slow:
     name: integration-tests-slow
@@ -94,6 +104,7 @@ jobs:
       GCP_US_EAST4_SERVER_URL: ${{ secrets.GCP_US_EAST4_SERVER_URL }}
       GCP_US_EAST4_SERVER_USERNAME: ${{ secrets.GCP_US_EAST4_SERVER_USERNAME }}
       GCP_US_EAST4_SERVER_PASSWORD: ${{ secrets.GCP_US_EAST4_SERVER_PASSWORD }}
+      PYTEST_RERUNS: ${{ inputs.reruns }}
     # TODO: add Windows testing for Python 3.11 and 3.12 back in
     # TODO: add macos testing back in
     if: ${{ ! startsWith(github.event.head_commit.message, 'GitBook:') && ! (inputs.os == 'windows-latest' && inputs.python-version == '3.11') && ! (inputs.os == 'windows-latest' && inputs.python-version == '3.12') && ! (inputs.os == 'macos-13' || inputs.os == 'macos-latest')  }}
diff --git a/.github/workflows/integration-test-slow.yml b/.github/workflows/integration-test-slow.yml
@@ -25,6 +25,11 @@ on:
         type: number
         required: false
         default: 30
+      reruns:
+        description: Pytest rerun count (0 disables)
+        type: number
+        required: false
+        default: 3
   workflow_dispatch:
     inputs:
       os:
@@ -71,6 +76,11 @@ on:
         type: number
         required: false
         default: 30
+      reruns:
+        description: Pytest rerun count (0 disables)
+        type: number
+        required: false
+        default: 3
 jobs:
   integration-tests-slow:
     name: integration-tests-slow
@@ -94,6 +104,7 @@ jobs:
       GCP_US_EAST4_SERVER_URL: ${{ secrets.GCP_US_EAST4_SERVER_URL }}
       GCP_US_EAST4_SERVER_USERNAME: ${{ secrets.GCP_US_EAST4_SERVER_USERNAME }}
       GCP_US_EAST4_SERVER_PASSWORD: ${{ secrets.GCP_US_EAST4_SERVER_PASSWORD }}
+      PYTEST_RERUNS: ${{ inputs.reruns }}
     # TODO: add Windows testing for Python 3.11 and 3.12 back in
     if: ${{ ! startsWith(github.event.head_commit.message, 'GitBook:') && ! (inputs.os == 'windows-latest' && inputs.python-version == '3.11') && ! (inputs.os == 'windows-latest' && inputs.python-version == '3.12') }}
     defaults:
diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml
@@ -26,6 +26,11 @@ on:
         type: string
         required: false
         default: ''
+      reruns:
+        description: Pytest rerun count (0 disables)
+        type: number
+        required: false
+        default: 3
   workflow_dispatch:
     inputs:
       os:
@@ -51,6 +56,11 @@ on:
         type: string
         required: false
         default: ''
+      reruns:
+        description: Pytest rerun count (0 disables)
+        type: number
+        required: false
+        default: 3
 jobs:
   unit-test:
     name: unit-test
@@ -61,6 +71,7 @@ jobs:
       PYTHONIOENCODING: utf-8
       UV_HTTP_TIMEOUT: 600
       OBJC_DISABLE_INITIALIZE_FORK_SAFETY: 'YES'
+      PYTEST_RERUNS: ${{ inputs.reruns }}
     if: ${{ ! startsWith(github.event.head_commit.message, 'GitBook:') && ! (inputs.os == 'windows-latest' && inputs.python-version == '3.11') && ! (inputs.os == 'windows-latest' && inputs.python-version == '3.12') }}
     defaults:
       run:
diff --git a/pyproject.toml b/pyproject.toml
@@ -270,6 +270,9 @@ max-complexity = 18
 # Use Google-style docstrings.
 convention = "google"
 
+[tool.bandit]
+skips = ["B615"]
+
 [tool.mypy]
 
 plugins = ["pydantic.mypy"]
diff --git a/scripts/check-security.sh b/scripts/check-security.sh
@@ -8,6 +8,6 @@ SRC=${1:-"src/zenml tests examples"}
 export ZENML_DEBUG=1
 export ZENML_ANALYTICS_OPT_IN=false
 
-bandit -r $SRC -ll \
+bandit -c pyproject.toml \
+    -r $SRC -ll \
     --exclude examples/llm_finetuning/scripts/prepare_alpaca.py
-
diff --git a/scripts/test-coverage-xml.sh b/scripts/test-coverage-xml.sh
@@ -14,28 +14,56 @@ TEST_SPLITS=${3:-"1"}
 TEST_GROUP=${4:-"1"}
 STORE_DURATIONS=${5:-""}
 
+# Control flaky test retries via environment variables.
+# - PYTEST_RERUNS: non-negative integer (0 disables reruns), defaults to 3 if unset or invalid.
+# - PYTEST_RERUNS_DELAY: non-negative integer delay between reruns in seconds, defaults to 5 if unset or invalid.
+# Validation guards against misconfiguration in CI while allowing explicit opt-out with 0.
+RERUNS_DEFAULT=3
+DELAY_DEFAULT=5
+
+if [[ -n "${PYTEST_RERUNS+x}" ]]; then
+    RERUNS="$PYTEST_RERUNS"
+else
+    RERUNS="$RERUNS_DEFAULT"
+fi
+if ! [[ "$RERUNS" =~ ^[0-9]+$ ]]; then
+    echo "Warning: PYTEST_RERUNS='$RERUNS' is invalid. Falling back to ${RERUNS_DEFAULT}." >&2
+    RERUNS="$RERUNS_DEFAULT"
+fi
+
+if [[ -n "${PYTEST_RERUNS_DELAY+x}" ]]; then
+    RERUNS_DELAY="$PYTEST_RERUNS_DELAY"
+else
+    RERUNS_DELAY="$DELAY_DEFAULT"
+fi
+if ! [[ "$RERUNS_DELAY" =~ ^[0-9]+$ ]]; then
+    echo "Warning: PYTEST_RERUNS_DELAY='$RERUNS_DELAY' is invalid. Falling back to ${DELAY_DEFAULT}." >&2
+    RERUNS_DELAY="$DELAY_DEFAULT"
+fi
+
+PYTEST_RERUN_ARGS=(--reruns "$RERUNS" --reruns-delay "$RERUNS_DELAY")
+
 export ZENML_DEBUG=1
 export ZENML_ANALYTICS_OPT_IN=false
 export EVIDENTLY_DISABLE_TELEMETRY=1
 
 ./zen-test environment provision $TEST_ENVIRONMENT
 
 # The '-vv' flag enables pytest-clarity output when tests fail.
-# Reruns a failing test 3 times with a 5 second delay.
 # Shows errors instantly in logs when test fails.
 if [ -n "$1" ]; then
     if [ "$STORE_DURATIONS" == "store-durations" ]; then
-        coverage run -m pytest $TEST_SRC --color=yes -vv --environment $TEST_ENVIRONMENT --no-provision --cleanup-docker --store-durations --durations-path=.test_durations --reruns 3 --reruns-delay 5 --instafail
+        coverage run -m pytest $TEST_SRC --color=yes -vv --environment $TEST_ENVIRONMENT --no-provision --cleanup-docker --store-durations --durations-path=.test_durations "${PYTEST_RERUN_ARGS[@]}" --instafail
     else
-        coverage run -m pytest $TEST_SRC --color=yes -vv --durations-path=.test_durations --splits=$TEST_SPLITS --group=$TEST_GROUP --splitting-algorithm least_duration --environment $TEST_ENVIRONMENT --no-provision --cleanup-docker --reruns 3 --reruns-delay 5 --instafail
+        coverage run -m pytest $TEST_SRC --color=yes -vv --durations-path=.test_durations --splits=$TEST_SPLITS --group=$TEST_GROUP --splitting-algorithm least_duration --environment $TEST_ENVIRONMENT --no-provision --cleanup-docker "${PYTEST_RERUN_ARGS[@]}" --instafail
     fi
 else
     if [ "$STORE_DURATIONS" == "store-durations" ]; then
-        coverage run -m pytest tests/unit --color=yes -vv --environment $TEST_ENVIRONMENT --no-provision --store-durations --durations-path=.test_durations --reruns 3 --reruns-delay 5 --instafail
-        coverage run -m pytest tests/integration --color=yes -vv --environment $TEST_ENVIRONMENT --no-provision --cleanup-docker --store-durations --durations-path=.test_durations --reruns 3 --reruns-delay 5 --instafail
+        coverage run -m pytest tests/unit --color=yes -vv --environment $TEST_ENVIRONMENT --no-provision --store-durations --durations-path=.test_durations "${PYTEST_RERUN_ARGS[@]}" --instafail
+        coverage run -m pytest tests/integration --color=yes -vv --environment $TEST_ENVIRONMENT --no-provision --cleanup-docker --store-durations --durations-path=.test_durations "${PYTEST_RERUN_ARGS[@]}" --instafail
     else
-        coverage run -m pytest tests/unit --color=yes -vv --durations-path=.test_durations --splits=$TEST_SPLITS --group=$TEST_GROUP --splitting-algorithm least_duration --environment $TEST_ENVIRONMENT --no-provision --reruns 3 --reruns-delay 5 --instafail
-        coverage run -m pytest tests/integration --color=yes -vv --durations-path=.test_durations --splits=$TEST_SPLITS --group=$TEST_GROUP --splitting-algorithm least_duration --environment $TEST_ENVIRONMENT --no-provision --cleanup-docker --reruns 3 --reruns-delay 5 --instafail
+        coverage run -m pytest tests/unit --color=yes -vv --durations-path=.test_durations --splits=$TEST_SPLITS --group=$TEST_GROUP --splitting-algorithm least_duration --environment $TEST_ENVIRONMENT --no-provision "${PYTEST_RERUN_ARGS[@]}" --instafail
+        coverage run -m pytest tests/integration --color=yes -vv --durations-path=.test_durations --splits=$TEST_SPLITS --group=$TEST_GROUP --splitting-algorithm least_duration --environment $TEST_ENVIRONMENT --no-provision --cleanup-docker "${PYTEST_RERUN_ARGS[@]}" --instafail
     fi
 fi
 
diff --git a/src/zenml/zen_stores/migrations/versions/502b4fa5fa88_adding_in_progress_to_runs.py b/src/zenml/zen_stores/migrations/versions/502b4fa5fa88_adding_in_progress_to_runs.py
@@ -37,7 +37,7 @@ def upgrade() -> None:
                 WHEN status IN ({finished_statuses}) THEN false 
                 ELSE true 
             END
-        """)
+        """)  # nosec
     )
 
     with op.batch_alter_table("pipeline_run", schema=None) as batch_op:

Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,7 @@ def upgrade() -> None:`
`37`	`37`	`WHEN status IN ({finished_statuses}) THEN false`
`38`	`38`	`ELSE true`
`39`	`39`	`END`
`40`		`- """)`
	`40`	`+ """) # nosec`
`41`	`41`	`)`
`42`	`42`
`43`	`43`	`with op.batch_alter_table("pipeline_run", schema=None) as batch_op:`