From e49ea2e3f53bf95802795ef281d3b12aef25e6c9 Mon Sep 17 00:00:00 2001
From: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
Date: Fri, 8 Nov 2024 02:31:42 +0100
Subject: [PATCH] Pytest runtime optimization

---
 .github/workflows/ci.yaml | 67 ++++++++++++++++++++++++++++++++++++++-
 .gitignore                |  3 +-
 script/split_tests.py     |  2 ++
 tests/conftest.py         | 29 +++++++++++++++++
 4 files changed, 99 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 549f819cd29ca3..ba183764f61496 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -40,6 +40,7 @@ env:
   CACHE_VERSION: 11-c1
   UV_CACHE_VERSION: 1
   MYPY_CACHE_VERSION: 9
+  PYTEST_RUNTIME_CACHE_VERSION: 1
   HA_SHORT_VERSION: "2025.1"
   DEFAULT_PYTHON: "3.12"
   ALL_PYTHON_VERSIONS: "['3.12', '3.13']"
@@ -934,6 +935,21 @@ jobs:
           key: >-
             ${{ runner.os }}-${{ steps.python.outputs.python-version }}-${{
             needs.info.outputs.python_cache_key }}
+      - name: Generate partial pytest runtime restore key
+        id: generate-pytest-runtime-key
+        run: |
+          echo "key=pytest-runtime-${{ env.PYTEST_RUNTIME_CACHE_VERSION }}-${{
+            env.HA_SHORT_VERSION }}-$(date -u '+%Y-%m-%dT%H:%M:%s')" >> $GITHUB_OUTPUT
+      - name: Restore pytest runtime cache
+        uses: actions/cache/restore@v4.1.2
+        with:
+          path: pytest-time-report-${{ env.DEFAULT_PYTHON }}.json
+          key: >-
+            ${{ runner.os }}-${{ env.DEFAULT_PYTHON }}-${{
+            steps.generate-pytest-runtime-key.outputs.key }}
+          restore-keys: |
+            ${{ runner.os }}-${{ env.DEFAULT_PYTHON }}-pytest-runtime-${{
+            env.PYTEST_RUNTIME_CACHE_VERSION }}-${{ env.HA_SHORT_VERSION }}-
       - name: Run split_tests.py
         run: |
           . venv/bin/activate
@@ -944,6 +960,12 @@ jobs:
           name: pytest_buckets
           path: pytest_buckets.txt
           overwrite: true
+      - name: Upload pytest collect output
+        uses: actions/upload-artifact@v4
+        with:
+          name: pytest-${{ github.run_number }}-collection
+          path: pytest-collect-output.txt
+          overwrite: true
 
   pytest-full:
     runs-on: ubuntu-24.04
@@ -1035,6 +1057,7 @@ jobs:
             ${cov_params[@]} \
             -o console_output_style=count \
             -p no:sugar \
+            --time-report-name pytest-time-report-${{ matrix.python-version }}-${{ matrix.group }}.json \
             $(sed -n "${{ matrix.group }},1p" pytest_buckets.txt) \
               2>&1 | tee pytest-${{ matrix.python-version }}-${{ matrix.group }}.txt
       - name: Upload pytest output
@@ -1042,7 +1065,9 @@ jobs:
         uses: actions/upload-artifact@v4.4.3
         with:
           name: pytest-${{ github.run_number }}-${{ matrix.python-version }}-${{ matrix.group }}
-          path: pytest-*.txt
+          path: |
+            pytest-*.txt
+            pytest-*.json
           overwrite: true
       - name: Upload coverage artifact
         if: needs.info.outputs.skip_coverage != 'true'
@@ -1057,6 +1082,45 @@ jobs:
         run: |
           ./script/check_dirty
 
+  pytest-time:
+    runs-on: ubuntu-24.04
+    needs:
+      - info
+      - pytest-full
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ${{ fromJSON(needs.info.outputs.python_versions) }}
+    name: Combine test runtimes
+    steps:
+      - name: Generate partial pytest runtime restore key
+        id: generate-pytest-runtime-key
+        run: |
+          echo "key=pytest-runtime-${{ env.PYTEST_RUNTIME_CACHE_VERSION }}-${{
+            env.HA_SHORT_VERSION }}-$(date -u '+%Y-%m-%dT%H:%M:%s')" >> $GITHUB_OUTPUT
+      - name: Download pytest runtime artifacts
+        uses: actions/download-artifact@v4.1.8
+        with:
+          pattern: pytest-${{ github.run_number }}-${{ matrix.python-version }}-*
+          merge-multiple: true
+      - name: Combine files into one
+        run: |
+          jq 'reduce inputs as $item ({}; . *= $item)' \
+            pytest-time-report-${{ matrix.python-version }}-*.json \
+            > pytest-time-report-${{ matrix.python-version }}.json
+      - name: Upload combined pytest runtime artifact
+        uses: actions/upload-artifact@v4.4.3
+        with:
+          name: pytest-${{ github.run_number }}-${{ matrix.python-version }}-time-report
+          path: pytest-time-report-${{ matrix.python-version }}.json
+      - name: Save pytest runtime cache
+        uses: actions/cache/save@v4.1.2
+        with:
+          path: pytest-time-report-${{ matrix.python-version }}.json
+          key: >-
+            ${{ runner.os }}-${{ matrix.python-version }}-${{
+            steps.generate-pytest-runtime-key.outputs.key }}
+
   pytest-mariadb:
     runs-on: ubuntu-24.04
     services:
@@ -1487,6 +1551,7 @@ jobs:
       - audit-licenses
       - installed-packages
       - pytest-full
+      - pytest-time
     timeout-minutes: 5
     steps:
       - name: Merge all pytest coverage artifacts
diff --git a/.gitignore b/.gitignore
index 241255253c5c4a..ee46f2976675e7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -136,4 +136,5 @@ tmp_cache
 .ropeproject
 
 # Will be created from script/split_tests.py
-pytest_buckets.txt
\ No newline at end of file
+pytest_buckets.txt
+pytest-time-report*.json
diff --git a/script/split_tests.py b/script/split_tests.py
index c64de46a0682c1..7b54388d82ee4d 100755
--- a/script/split_tests.py
+++ b/script/split_tests.py
@@ -181,6 +181,8 @@ def collect_tests(path: Path) -> TestFolder:
 
     folder = TestFolder(path)
 
+    Path("pytest-collect-output.txt").write_text(result.stdout)
+
     for line in result.stdout.splitlines():
         if not line.strip():
             continue
diff --git a/tests/conftest.py b/tests/conftest.py
index 2cefe72f41487b..f991ed542a43f6 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -9,8 +9,10 @@
 import functools
 import gc
 import itertools
+import json
 import logging
 import os
+from pathlib import Path
 import reprlib
 from shutil import rmtree
 import sqlite3
@@ -19,6 +21,7 @@
 from typing import TYPE_CHECKING, Any, cast
 from unittest.mock import AsyncMock, MagicMock, Mock, _patch, patch
 
+from _pytest.terminal import TerminalReporter
 from aiohttp import client
 from aiohttp.test_utils import (
     BaseTestServer,
@@ -140,6 +143,30 @@
 def pytest_addoption(parser: pytest.Parser) -> None:
     """Register custom pytest options."""
     parser.addoption("--dburl", action="store", default="sqlite://")
+    parser.addoption(
+        "--time-report-name", action="store", default="pytest-time-report.json"
+    )
+
+
+class PytestExecutionTimer:  # noqa: D101
+    def pytest_terminal_summary(  # noqa: D102
+        self,
+        terminalreporter: TerminalReporter,
+        exitstatus: pytest.ExitCode,
+        config: pytest.Config,
+    ) -> None:
+        if config.option.collectonly:
+            return
+        raw_data: dict[str, list[float]] = {}
+        for replist in terminalreporter.stats.values():
+            for rep in replist:
+                if isinstance(rep, pytest.TestReport):
+                    raw_data.setdefault(rep.location[0], []).append(rep.duration)
+        data = {filename: sum(values) for filename, values in raw_data.items()}
+        time_report_filename = config.option.time_report_name
+        file = Path(__file__).parents[1].joinpath(time_report_filename)
+        with open(file, "w", encoding="utf-8") as fp:
+            json.dump(data, fp, indent=2)
 
 
 def pytest_configure(config: pytest.Config) -> None:
@@ -155,6 +182,8 @@ def pytest_configure(config: pytest.Config) -> None:
     # See https://github.com/syrupy-project/syrupy/pull/901
     SnapshotSession.finish = override_syrupy_finish
 
+    config.pluginmanager.register(PytestExecutionTimer())
+
 
 def pytest_runtest_setup() -> None:
     """Prepare pytest_socket and freezegun.