feat: support Schemathesis hooks

aaltat · aaltat · commit 34cc699fd9f3 · 2026-01-11T21:51:10.000+02:00
diff --git a/README.md b/README.md
@@ -107,3 +107,55 @@ class AuthExtension:
 Then with all API calls, will have
 [basic auth](https://en.wikipedia.org/wiki/Basic_access_authentication) set in the
 headers for all calls made to your API endpoint.
+
+## Schemathesis hook support
+Library supports extending Schemathesis by defining hooks. Hooks allows users to customize how
+Schemathesis generates test data, validates responses, and handles requests through hooks, custom
+checks, and data generation strategies. For more details about Schemathesis hooks, refer to
+Schemathesis extending documentation: https://schemathesis.readthedocs.io/en/stable/guides/extending/
+
+Example if there need to add custom header in each request, then it is possible to import library
+with hook:
+
+```robotframework
+*** Settings ***
+Variables           authentication.py
+Library             SchemathesisLibrary
+...                     url=http://127.0.0.1/openapi.json
+...                     hook=${CURDIR}/hook_filter.py
+
+Test Template       Wrapper
+
+
+*** Test Cases ***
+All Tests
+    Wrapper    test_case_1
+
+
+*** Keywords ***
+Wrapper
+    [Arguments]    ${case}
+    Call And Validate    ${case}    auth=${BASIC_AUTH_TUPLE}
+
+```
+
+And when hook_filter.py looks like:
+
+```python
+import schemathesis
+
+
+global_count_count = 0
+
+
+@schemathesis.hook
+def filter_query(ctx, query) -> bool:
+    method = ctx.operation.method.lower().strip()
+    if method == "put":
+        global global_count_count
+        global_count_count += 1
+        if global_count_count > 2:
+            return False
+    return True
+```
+Then only two test with `PUT` request are generated.
diff --git a/atest/library/hook_filter.py b/atest/library/hook_filter.py
@@ -0,0 +1,15 @@
+import schemathesis
+
+
+global_count_count = 0
+
+
+@schemathesis.hook
+def filter_query(ctx, query) -> bool:
+    method = ctx.operation.method.lower().strip()
+    if method == "put":
+        global global_count_count
+        global_count_count += 1
+        if global_count_count > 2:
+            return False
+    return True
diff --git a/atest/library/hook_flatmap.py b/atest/library/hook_flatmap.py
@@ -0,0 +1,9 @@
+import schemathesis
+
+
+@schemathesis.hook
+def map_headers(ctx, headers):
+    if headers is None:
+        headers = {}
+    headers["flatmap-headers"] = "I am here too"
+    return headers
diff --git a/atest/library/hook_map.py b/atest/library/hook_map.py
@@ -0,0 +1,9 @@
+import schemathesis
+
+
+@schemathesis.hook
+def map_headers(ctx, headers):
+    if headers is None:
+        headers = {}
+    headers["map-headers"] = "I am here"
+    return headers
diff --git a/atest/library/hooks_filter.robot b/atest/library/hooks_filter.robot
@@ -0,0 +1,19 @@
+*** Settings ***
+Variables           authentication.py
+Library             SchemathesisLibrary
+...                     url=http://127.0.0.1/openapi.json
+...                     hook=${CURDIR}/hook_filter.py
+
+Test Template       Wrapper
+
+
+*** Test Cases ***
+All Tests
+    Wrapper    test_case_1
+
+
+*** Keywords ***
+Wrapper
+    [Arguments]    ${case}
+    ${r} =    Call And Validate    ${case}    auth=${BASIC_AUTH_TUPLE}
+    Log    ${r.json()}
diff --git a/atest/library/hooks_map.robot b/atest/library/hooks_map.robot
@@ -0,0 +1,20 @@
+*** Settings ***
+Variables           authentication.py
+Library             SchemathesisLibrary
+...                     url=http://127.0.0.1/openapi.json
+...                     hook=${CURDIR}/hook_map.py
+...                     max_examples=4
+
+Test Template       Wrapper
+
+
+*** Test Cases ***
+All Tests
+    Wrapper    test_case_1
+
+
+*** Keywords ***
+Wrapper
+    [Arguments]    ${case}
+    ${r} =    Call And Validate    ${case}    auth=${BASIC_AUTH_TUPLE}
+    Log    ${r.json()}
diff --git a/atest/library/hooks_map_and_flatmap.robot b/atest/library/hooks_map_and_flatmap.robot
@@ -0,0 +1,20 @@
+*** Settings ***
+Variables           authentication.py
+Library             SchemathesisLibrary
+...                     url=http://127.0.0.1/openapi.json
+...                     hook=${CURDIR}/hook_map.py;${CURDIR}/hook_flatmap.py
+...                     max_examples=4
+
+Test Template       Wrapper
+
+
+*** Test Cases ***
+All Tests
+    Wrapper    test_case_1
+
+
+*** Keywords ***
+Wrapper
+    [Arguments]    ${case}
+    ${r} =    Call And Validate    ${case}    auth=${BASIC_AUTH_TUPLE}
+    Log    ${r.json()}
diff --git a/atest/test/TestSuiteChecker.py b/atest/test/TestSuiteChecker.py
@@ -1,6 +1,7 @@
 from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
+from typing import TypedDict
 
 from robot.api import ExecutionResult, ResultVisitor, logger
 from robot.libraries.BuiltIn import BuiltIn
@@ -78,6 +79,44 @@ def _normalize_kw_name(self, name: str) -> str:
         return name.replace(" ", "").replace("_", "").lower()
 
 
+class ExecutionCountChecker(ResultVisitor):
+    def __init__(self):
+        self.test_names_counts = defaultdict(int)
+        self.test_count = 0
+
+    def visit_test(self, test: TestCase):
+        self.test_count += 1
+        test_name = test.name.split(" ")[0]
+        self.test_names_counts[test_name] += 1
+        return super().visit_test(test)
+
+
+class ExecutionLogChecker(ResultVisitor):
+    def __init__(self, match_start: str, kw_name: str = "Call And Validate"):
+        self.match_start = match_start
+        self.kw_name = kw_name.lower()
+        self.test_logs = {}
+
+    def visit_message(self, message: Message):
+        if not message.message.startswith(self.match_start):
+            return super().visit_message(message)
+        if not (
+            message.parent
+            and isinstance(message.parent, Keyword)
+            and message.parent.name.lower() == self.kw_name
+        ):
+            return super().visit_message(message)
+        parent = message.parent
+        while not isinstance(parent, TestCase):
+            if not parent.parent:
+                return super().visit_message(message)
+            parent = parent.parent
+        test = parent
+        logger.debug(f"test case: {test}, message: {message.message}")
+        self.test_logs[test.name] = message.message
+        return super().visit_message(message)
+
+
 class TestSuiteChecker:
     def check_suite(self, output_xml: str, test_description: list[TestCaseData]):
         """Check the test suite for test cases names."""
@@ -108,6 +147,55 @@ def check_suite(self, output_xml: str, test_description: list[TestCaseData]):
             self._check_test_logs(test.name, test.logs, checker.received_test_logs, should_match)
         logger.info(f"Checked test suite in {output_xml}")
 
+    def check_test_names_and_counts(
+        self, output_xml: str, min_count: int, max_count: int, expected_names_and_counts: dict[str, int]
+    ):
+        """Check the test case start with correct names and has correct number of test cases."""
+        result = ExecutionResult(output_xml)
+        checker = ExecutionCountChecker()
+        result.visit(checker)
+        result.save(output_xml)
+        logger.info(f"Checked test suite in {output_xml}")
+        assert checker.test_count >= min_count, (
+            f"Expected at least {min_count} tests, but found {checker.test_count}."
+        )
+        assert checker.test_count <= max_count, (
+            f"Expected at most {max_count} tests, but found {checker.test_count}."
+        )
+        for test_name, count in expected_names_and_counts.items():
+            assert test_name in checker.test_names_counts, (
+                f"Expected test name {test_name} not found in {checker.test_names_counts.keys()} "
+            )
+            assert checker.test_names_counts[test_name] >= count, (
+                f"Expected {count} for test name {test_name} but it was {checker.test_names_counts[test_name]}"
+            )
+
+    def check_specific_test_logs(
+        self,
+        output_xml: str,
+        log_message_start: str,
+        kw_name: str = "Call And Validate",
+        number_of_tests: int = 15,
+        string_in_log: None | str = None,
+    ):
+        assert number_of_tests > 0, "number_of_tests must be greater than 0"
+        assert string_in_log, "string_in_log must be provided"
+        result = ExecutionResult(output_xml)
+        checker = ExecutionLogChecker(log_message_start, kw_name)
+        result.visit(checker)
+        result.save(output_xml)
+        logger.info(f"Checked test suite in {output_xml}")
+        assert len(checker.test_logs) == number_of_tests, (
+            f"Expected {number_of_tests} tests with logs starting with '{log_message_start}', "
+            f"but found {len(checker.test_logs)}."
+        )
+        for test, log in checker.test_logs.items():
+            logger.info(f"Checking log for test '{test}': {log}")
+            assert string_in_log in log, (
+                f"Expected log to contain '{string_in_log}', but it was not found in log: {log} in test '{test}'."
+            )
+        logger.debug(f"Found logs for tests: {checker.test_logs}")
+
     def _check_test_logs(
         self,
         test_name: str,
diff --git a/atest/test/all_cases.resource b/atest/test/all_cases.resource
@@ -93,3 +93,20 @@ Get All Cases Test Description With Set Headers
         END
     END
     RETURN    ${test_description}
+
+Get All Cases Test Description With Map Headers
+    [Arguments]    ${count: int}=${5}
+    ${test_description} =    Get All Cases Test Description    ${count}
+    FOR    ${description}    IN    @{test_description}
+        VAR    ${logs} =    ${description}[logs]
+        FOR    ${i}    ${log}    IN ENUMERATE    @{logs}
+            IF    $log.startswith("Case headers")
+                Log    Original log: ${log}
+                VAR    ${new_log} =
+                ...    Case headers {'map-headers': 'I am here'} body 'Not set' cookies {} path parameters {}
+                Log    Modified log: ${new_log}
+                Set List Value    ${logs}    ${i}    ${new_log}
+            END
+        END
+    END
+    RETURN    ${test_description}
diff --git a/atest/test/hooks_filter.robot b/atest/test/hooks_filter.robot
@@ -0,0 +1,11 @@
+*** Settings ***
+Resource        runner.resource
+Resource        all_cases.resource
+
+Suite Setup     Run Suite
+
+
+*** Test Cases ***
+Check All Cases
+    VAR    &{tests} =    GET=11    PUT=1
+    Check Test Names And Counts    ${LIBRARY_OUTPUT_XML}    17    18    ${tests}
diff --git a/atest/test/hooks_map.robot b/atest/test/hooks_map.robot
@@ -0,0 +1,19 @@
+*** Settings ***
+Resource        runner.resource
+Resource        all_cases.resource
+
+Suite Setup     Run Suite
+
+
+*** Test Cases ***
+Check All Cases
+    Check Specific Test Logs
+    ...    ${LIBRARY_OUTPUT_XML}
+    ...    Case headers
+    ...    number_of_tests=17
+    ...    string_in_log='map-headers': 'I am here'
+    Check Specific Test Logs
+    ...    ${LIBRARY_OUTPUT_XML}
+    ...    Request:
+    ...    number_of_tests=17
+    ...    string_in_log='map-headers': 'I am here'
diff --git a/atest/test/hooks_map_and_flatmap.robot b/atest/test/hooks_map_and_flatmap.robot
@@ -0,0 +1,29 @@
+*** Settings ***
+Resource        runner.resource
+Resource        all_cases.resource
+
+Suite Setup     Run Suite
+
+
+*** Test Cases ***
+Check All Cases
+    Check Specific Test Logs
+    ...    ${LIBRARY_OUTPUT_XML}
+    ...    Case headers
+    ...    number_of_tests=17
+    ...    string_in_log='map-headers': 'I am here'
+    Check Specific Test Logs
+    ...    ${LIBRARY_OUTPUT_XML}
+    ...    Case headers
+    ...    number_of_tests=17
+    ...    string_in_log='flatmap-headers': 'I am here too'
+    Check Specific Test Logs
+    ...    ${LIBRARY_OUTPUT_XML}
+    ...    Request:
+    ...    number_of_tests=17
+    ...    string_in_log='map-headers': 'I am here'
+    Check Specific Test Logs
+    ...    ${LIBRARY_OUTPUT_XML}
+    ...    Request:
+    ...    number_of_tests=17
+    ...    string_in_log='flatmap-headers': 'I am here too'
diff --git a/src/SchemathesisLibrary/__init__.py b/src/SchemathesisLibrary/__init__.py
@@ -71,14 +71,15 @@ class SchemathesisLibrary(DynamicCore):
     ROBOT_LISTENER_API_VERSION = 3
     ROBOT_LIBRARY_SCOPE = "TEST SUITE"
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         *,
         headers: "dict[str, Any]|None" = None,
         max_examples: int = 5,
         path: "Path|None" = None,
         url: "str|None" = None,
         auth: str | None = None,
+        hook: str | None = None,
     ) -> None:
         """The SchemathesisLibrary can be initialized with the following arguments:
 
@@ -88,6 +89,7 @@ def __init__(
         | ``path`` | Path to the OpenAPI schema file. Using either ``path`` or ``url`` is mandatory. |
         | ``url`` | URL where the OpenAPI schema can be downloaded. |
         | ``auth`` | Optional authentication class to be used passed for Schemathesis authentication when test cases are executed. |
+        | ``hook`` | Optional Schemathesis hook https://schemathesis.readthedocs.io/en/stable/guides/extending/. |
 
         The ``headers`` argument is only needed when the schema is downloaded from a URL and there is need to pass example
         authentication headers to the endpoint. The ``headers`` is not used in the API calls are made during test execution.
@@ -104,10 +106,21 @@ def __init__(
         See example from
         [https://github.com/aaltat/robotframework-schemathesis?tab=readme-ov-file##dynamic-token-authentication|README.md]
         file
+
+        The ``hook`` argument can be used to specify Schemathesis hooks to modify the generated test cases. The hooks
+        implementation must follow the Schemathesis documentation:
+        https://schemathesis.readthedocs.io/en/stable/guides/extending/ Hooks importing follows same rules as importing
+        [https://robotframework.org/robotframework/latest/RobotFrameworkUserGuide.html#specifying-library-to-import|test libraries]
+        in Robot Framework. Multiple hooks can be specified by separating them with semicolon (;).
         """
         self.ROBOT_LIBRARY_LISTENER = self
         SchemathesisReader.options = Options(
-            headers=headers, max_examples=max_examples, path=path, url=url, auth=auth
+            headers=headers,
+            max_examples=max_examples,
+            path=path,
+            url=url,
+            auth=auth,
+            hook=hook,
         )
         self.data_driver = DataDriver(reader_class=SchemathesisReader)
         DynamicCore.__init__(self, [])
diff --git a/src/SchemathesisLibrary/schemathesisreader.py b/src/SchemathesisLibrary/schemathesisreader.py
