🧪 Add --firecrest-requests pytest option (#25)

.github/workflows/tests.yml

@@ -49,7 +49,7 @@ jobs:
         python -m pip install --upgrade pip
         pip install -e .[dev]
     - name: Test with pytest
-      run: pytest -vv --cov=aiida_firecrest --cov-report=xml --cov-report=term-missing
+      run: pytest -vv --firecrest-requests --cov=aiida_firecrest --cov-report=xml --cov-report=term
 
     - name: Upload coverage reports to Codecov
       uses: codecov/codecov-action@v3

README.md

@@ -167,6 +167,13 @@ tox -- --firecrest-config=".firecrest-demo-config.json" --firecrest-no-clean
 See [firecrest_demo.py](firecrest_demo.py) for how to start up a demo server,
 and also [server-tests.yml](.github/workflows/server-tests.yml) for how the tests are run against the demo server on GitHub Actions.
 
+If you want to analyse statistics of the API requests made by each test,
+you can use the `--firecrest-requests` option:
+
+```bash
+tox -- --firecrest-requests
+```
+
 ### Notes on using the demo server on MacOS
 
 A few issues have been noted when using the demo server on MacOS (non-Mx):

aiida_firecrest/utils_test.py

@@ -83,6 +83,7 @@ def mock_request(
         files: dict[str, tuple[str, BinaryIO]] | None = None,
         **kwargs: Any,
     ) -> Response:
+        """Mock a request to the Firecrest server."""
         response = Response()
         response.encoding = "utf-8"
         response.url = url if isinstance(url, str) else url.decode("utf-8")

conftest.py

@@ -1,6 +1,4 @@
-# add option to pytest
-# note this file must be at the root of the project
-
+"""Pytest configuration that must be at the root level."""
 pytest_plugins = ["aiida.manage.tests.pytest_fixtures"]
 
 
@@ -13,16 +11,8 @@ def pytest_addoption(parser):
         action="store_true",
         help="Don't clean up server after tests (for debugging)",
     )
-
-
-def pytest_report_header(config):
-    if config.getoption("--firecrest-config"):
-        header = [
-            "Running against FirecREST server: {}".format(
-                config.getoption("--firecrest-config")
-            )
-        ]
-        if config.getoption("--firecrest-no-clean"):
-            header.append("Not cleaning up FirecREST server after tests!")
-        return header
-    return ["Running against Mock FirecREST server"]
+    parser.addoption(
+        "--firecrest-requests",
+        action="store_true",
+        help="Collect and print telemetry data for API requests",
+    )

tests/conftest.py

@@ -1,20 +1,66 @@
+"""Pytest configuration for the aiida-firecrest tests.
+
+This sets up the Firecrest server to use, and telemetry for API requests.
+"""
 from __future__ import annotations
 
+from functools import partial
 from json import load as json_load
 from pathlib import Path
+from typing import Any, Callable
+from urllib.parse import urlparse
 
+from _pytest.terminal import TerminalReporter
 import firecrest as f7t
 import pytest
 import requests
+import yaml
 
 from aiida_firecrest.utils_test import FirecrestConfig, FirecrestMockServer
 
 
+def pytest_report_header(config):
+    if config.getoption("--firecrest-config"):
+        header = [
+            "Running against FirecREST server: {}".format(
+                config.getoption("--firecrest-config")
+            )
+        ]
+        if config.getoption("--firecrest-no-clean"):
+            header.append("Not cleaning up FirecREST server after tests!")
+        return header
+    return ["Running against Mock FirecREST server"]
+
+
+def pytest_terminal_summary(
+    terminalreporter: TerminalReporter, exitstatus: int, config: pytest.Config
+):
+    """Called after all tests have run."""
+    data = config.stash.get("firecrest_requests", None)
+    if data is None:
+        return
+    terminalreporter.write(
+        yaml.dump(
+            {"Firecrest requests telemetry": data},
+            default_flow_style=False,
+            sort_keys=True,
+        )
+    )
+
+
 @pytest.fixture(scope="function")
-def firecrest_server(request, monkeypatch, tmp_path: Path):
+def firecrest_server(
+    pytestconfig: pytest.Config,
+    request: pytest.FixtureRequest,
+    monkeypatch,
+    tmp_path: Path,
+):
     """A fixture which provides a mock Firecrest server to test against."""
-    config_path = request.config.getoption("--firecrest-config")
-    no_clean = request.config.getoption("--firecrest-no-clean")
+    config_path: str | None = request.config.getoption("--firecrest-config")
+    no_clean: bool = request.config.getoption("--firecrest-no-clean")
+    record_requests: bool = request.config.getoption("--firecrest-requests")
+    telemetry: RequestTelemetry | None = None
+
     if config_path is not None:
         # if given, use this config
         with open(config_path, encoding="utf8") as handle:
@@ -33,6 +79,18 @@ def firecrest_server(request, monkeypatch, tmp_path: Path):
             ),
         )
         client.mkdir(config.machine, config.scratch_path, p=True)
+
+        if record_requests:
+            telemetry = RequestTelemetry()
+            monkeypatch.setattr(requests, "get", partial(telemetry.wrap, requests.get))
+            monkeypatch.setattr(
+                requests, "post", partial(telemetry.wrap, requests.post)
+            )
+            monkeypatch.setattr(requests, "put", partial(telemetry.wrap, requests.put))
+            monkeypatch.setattr(
+                requests, "delete", partial(telemetry.wrap, requests.delete)
+            )
+
         yield config
         # Note this shouldn't really work, for folders but it does :shrug:
         # because they use `rm -r`:
@@ -42,8 +100,39 @@ def firecrest_server(request, monkeypatch, tmp_path: Path):
     else:
         # otherwise use mock server
         server = FirecrestMockServer(tmp_path)
-        monkeypatch.setattr(requests, "get", server.mock_request)
-        monkeypatch.setattr(requests, "post", server.mock_request)
-        monkeypatch.setattr(requests, "put", server.mock_request)
-        monkeypatch.setattr(requests, "delete", server.mock_request)
+        if record_requests:
+            telemetry = RequestTelemetry()
+            mock_request = partial(telemetry.wrap, server.mock_request)
+        else:
+            mock_request = server.mock_request
+        monkeypatch.setattr(requests, "get", mock_request)
+        monkeypatch.setattr(requests, "post", mock_request)
+        monkeypatch.setattr(requests, "put", mock_request)
+        monkeypatch.setattr(requests, "delete", mock_request)
         yield server.config
+
+    # save data on the server
+    if telemetry is not None:
+        test_name = request.node.name
+        pytestconfig.stash.setdefault("firecrest_requests", {})[
+            test_name
+        ] = telemetry.counts
+
+
+class RequestTelemetry:
+    """A to gather telemetry on requests."""
+
+    def __init__(self) -> None:
+        self.counts = {}
+
+    def wrap(
+        self,
+        method: Callable[..., requests.Response],
+        url: str | bytes,
+        **kwargs: Any,
+    ) -> requests.Response:
+        """Wrap a requests method to gather telemetry."""
+        endpoint = urlparse(url if isinstance(url, str) else url.decode("utf-8")).path
+        self.counts.setdefault(endpoint, 0)
+        self.counts[endpoint] += 1
+        return method(url, **kwargs)

tests/test_calculation.py

@@ -42,6 +42,7 @@ def _firecrest_computer(firecrest_server: FirecrestConfig):
 @pytest.fixture(name="no_retries")
 def _no_retries():
     """Remove calcjob retries, to make failing the test faster."""
+    # TODO calculation seems to hang on errors still
     max_attempts = manage.get_config().get_option(MAX_ATTEMPTS_OPTION)
     manage.get_config().set_option(MAX_ATTEMPTS_OPTION, 1)
     yield
@@ -67,10 +68,6 @@ def test_calculation_basic(firecrest_computer: orm.Computer):
     # https://github.com/eth-cscs/firecrest/issues/191
     builder.metadata.options.submit_script_filename = "aiidasubmit.sh"
 
-    # TODO reset in fixture? also the calculation seems to hang on errors still
-    # rather than failing immediately
-    manage.get_config().set_option(MAX_ATTEMPTS_OPTION, 1)
-
     _, node = engine.run_get_node(builder)
     assert node.is_finished_ok