From c34f2de973a42e64ea6bec4f0d4e2c1142250212 Mon Sep 17 00:00:00 2001 From: "A.Shpak" Date: Tue, 1 Apr 2025 14:51:14 +0300 Subject: [PATCH] feat: add cache support, timeouts and info field for probes --- .github/workflows/tests.yml | 23 +++- .gitignore | 1 + LICENSE | 2 +- README.md | 259 ++++++++++++++++++++++++++++-------- justfile | 2 +- probirka/_probes.py | 115 ++++++++++++---- probirka/_probirka.py | 53 +++++--- probirka/_results.py | 4 +- pyproject.toml | 1 + tests/conftest.py | 4 + tests/test_aiohttp.py | 144 ++++++++++++++++++++ tests/test_fastapi.py | 132 ++++++++++++++++++ tests/test_probes.py | 224 ++++++++++++++++++++++++++++++- tests/test_probirka.py | 49 ++++++- uv.lock | 46 ++++++- 15 files changed, 946 insertions(+), 113 deletions(-) create mode 100644 tests/test_aiohttp.py create mode 100644 tests/test_fastapi.py diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index a8dce96..72e2eae 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -6,13 +6,26 @@ on: branches: [ main ] jobs: - tests: + linters: runs-on: ubuntu-latest strategy: matrix: just-trigger: + - "fmt" - "lint" - - "tests" + - "mypy" + steps: + - uses: actions/checkout@v4 + - uses: extractions/setup-just@v2 + - uses: astral-sh/setup-uv@v4 + - run: uv python install 3.8 + - run: uv sync --all-extras --dev + - run: just ${{ matrix.just-trigger }} + + tests: + runs-on: ubuntu-latest + strategy: + matrix: python-version: - "3.8" - "3.9" @@ -26,4 +39,8 @@ jobs: - uses: astral-sh/setup-uv@v4 - run: uv python install ${{ matrix.python-version }} - run: uv sync --all-extras --dev - - run: just ${{ matrix.just-trigger }} + - run: just tests + - uses: coverallsapp/github-action@master + with: + github-token: ${{ secrets.github_token }} + path-to-lcov: tests.lcov \ No newline at end of file diff --git a/.gitignore b/.gitignore index 923f7dc..cf88a12 100644 --- a/.gitignore +++ b/.gitignore @@ -44,6 +44,7 @@ htmlcov/ .nox/ .coverage .coverage.* +*.lcov .cache nosetests.xml coverage.xml diff --git a/LICENSE b/LICENSE index 5008eb4..8bb0b6c 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2024 KODE LLC +Copyright (c) 2025 KODE LLC Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index 54b5121..b1c775c 100644 --- a/README.md +++ b/README.md @@ -5,19 +5,11 @@ Python 3 library to write simple asynchronous health checks (probes). [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff) [![PyPI](https://img.shields.io/pypi/v/probirka.svg)](https://pypi.python.org/pypi/probirka) [![PyPI](https://img.shields.io/pypi/dm/probirka.svg)](https://pypi.python.org/pypi/probirka) +[![Coverage Status](https://coveralls.io/repos/github/appKODE/probirka/badge.svg?branch=polish-docs)](https://coveralls.io/github/appKODE/probirka?branch=polish-docs) ## Overview -Probirka is a Python library designed to facilitate the creation of simple asynchronous health checks, also known as probes. It allows you to define custom probes to monitor the health of various services and components in your application. - -## Features - -- Simple API for defining asynchronous health checks -- Support for custom probes -- Integration with FastAPI and aiohttp -- Ability to add custom information to health check results -- Grouping of probes for selective execution -- Timeout support for probes +Probirka is a lightweight and flexible Python library for implementing asynchronous health checks in your applications. It provides a simple yet powerful API for monitoring the health of various components and services, making it ideal for microservices architectures, containerized applications, and distributed systems. ## Installation @@ -27,97 +19,260 @@ Install Probirka using pip: pip install probirka ``` -## Usage +## Quick Start -Here is a simple example of how to use Probirka to create an asynchronous health check: +Here is a simple example of how to use Probirka to create health checks using decorators: ```python import asyncio -from probirka import Probe, HealthCheck +from probirka import Probirka + +# Create a Probirka instance +probirka = Probirka() + +# Add some custom information +probirka.add_info("version", "1.0.0") +probirka.add_info("environment", "production") + +# Define health checks using decorators +@probirka.add(name="database") # This probe will always run +async def check_database(): + # Simulate a database check + await asyncio.sleep(1) + return True + +@probirka.add(groups=["cache"]) # This probe will only run when cache group is requested +async def check_cache(): + # Simulate a cache check + await asyncio.sleep(1) + return False # Simulate a failed check -class DatabaseProbe(Probe): - async def check(self): +@probirka.add(groups=["external"]) # This probe will only run when external group is requested +def check_external_service(): + # Synchronous check example + return True + +async def main(): + # Run only required probes (without groups) + basic_results = await probirka.run() + print("Basic check results:", basic_results) + + # Run required probes + cache group probes + cache_results = await probirka.run(with_groups=["cache"]) + print("Cache check results:", cache_results) + + # Run required probes + multiple groups + full_results = await probirka.run(with_groups=["cache", "external"]) + print("Full check results:", full_results) + +if __name__ == "__main__": + asyncio.run(main()) +``` + +Alternatively, you can create custom probes by inheriting from the `ProbeBase` class: + +```python +from probirka import Probirka, ProbeBase +import asyncio + +class DatabaseProbe(ProbeBase): + async def _check(self): # Simulate a database check await asyncio.sleep(1) return True -class CacheProbe(Probe): - async def check(self): +class CacheProbe(ProbeBase): + async def _check(self): # Simulate a cache check await asyncio.sleep(1) return False # Simulate a failed check async def main(): - health_check = HealthCheck(probes=[DatabaseProbe(), CacheProbe()]) - health_check.add_info("version", "1.0.0") - health_check.add_info("environment", "production") - results = await health_check.run() + probirka = Probirka(probes=[DatabaseProbe(), CacheProbe()]) + probirka.add_info("version", "1.0.0") + probirka.add_info("environment", "production") + results = await probirka.run() print(results) if __name__ == "__main__": asyncio.run(main()) ``` -This example defines two probes, `DatabaseProbe` and `CacheProbe`, and runs them as part of a health check. The `CacheProbe` simulates a failed check. Additionally, it adds custom user data using the `add_info` method. +## Advanced Usage + +### Creating Custom Probes -Example output: +You can create custom probes by inheriting from the `ProbeBase` class: ```python -HealthCheckResult( - ok=False, - info={'version': '1.0.0', 'environment': 'production'}, - started_at=datetime.datetime(2023, 10, 10, 10, 0, 0), - total_elapsed=datetime.timedelta(seconds=1), - checks=[ - ProbeResult(name='DatabaseProbe', ok=True, elapsed=datetime.timedelta(seconds=1)), - ProbeResult(name='CacheProbe', ok=False, elapsed=datetime.timedelta(seconds=1)) - ] -) +from probirka import ProbeBase +import asyncio + +class CustomProbe(ProbeBase): + def __init__(self, name="CustomProbe"): + super().__init__(name=name) + + async def _check(self): + # Implement your health check logic here + return True ``` -## Integration with FastAPI +### Adding Metadata to Probes -You can integrate Probirka with FastAPI as follows: +You can add metadata to your probes: + +```python +from probirka import ProbeBase +import asyncio + +class DatabaseProbe(ProbeBase): + async def _check(self): + await asyncio.sleep(1) + self.add_info("connection_pool_size", 10) + self.add_info("active_connections", 5) + return True +``` + +The added information will be included in the probe results and can be accessed through the `info` field of each probe result. This is useful for providing additional context about the probe's state or performance metrics. + +### Grouping Probes + +Probes can be organized into required and optional groups. Probes without groups are always executed, while probes with groups are only executed when explicitly requested: + +```python +import asyncio +from probirka import Probirka + +# Create a Probirka instance +probirka = Probirka() + +# Required probe (will always run) +@probirka.add(name="database") +async def check_database(): + await asyncio.sleep(1) + return True + +# Optional probes (will only run when their groups are requested) +@probirka.add(groups=["cache"]) +async def check_cache(): + await asyncio.sleep(1) + return True + +@probirka.add(groups=["external"]) +async def check_external_service(): + return True + +async def main(): + # Run only required probes (database) + basic_results = await probirka.run() + print("Basic check results:", basic_results) + + # Run required probes + cache group + cache_results = await probirka.run(with_groups=["cache"]) + print("Cache check results:", cache_results) + + # Run required probes + multiple groups + full_results = await probirka.run(with_groups=["cache", "external"]) + print("Full check results:", full_results) + +if __name__ == "__main__": + asyncio.run(main()) +``` + +### Setting Timeouts + +You can set timeouts for individual probes: + +```python +from probirka import ProbeBase +import asyncio + +class SlowProbe(ProbeBase): + async def _check(self): + await asyncio.sleep(2) # This will timeout + return True + +probe = SlowProbe(timeout=1.0) # 1 second timeout +``` + +### Caching Results + +```python +from typing import Optional +from probirka import Probirka, ProbeBase +import asyncio + +# Create a Probirka instance with global caching settings +probirka = Probirka(success_ttl=60, failed_ttl=10) # Cache successful results for 60s, failed for 10s + +# Add a probe with custom caching settings +@probirka.add(success_ttl=300) # Cache successful results for 5 minutes +async def check_database(): + # Simulate a database check + await asyncio.sleep(1) + return True + +# Or create a custom probe with caching +class DatabaseProbe(ProbeBase): + def __init__(self, success_ttl: Optional[int] = None, failed_ttl: Optional[int] = None): + super().__init__(success_ttl=success_ttl, failed_ttl=failed_ttl) + + async def _check(self) -> bool: + # Simulate a database check + await asyncio.sleep(1) + return True +``` + +The caching mechanism works as follows: +- If `success_ttl` is set, successful results will be cached for the specified number of seconds +- If `failed_ttl` is set, failed results will be cached for the specified number of seconds +- If both are set to `None` (default), no caching will be performed +- Global settings in `Probirka` instance can be overridden by individual probe settings + +## Integration Examples + +### FastAPI Integration ```python from fastapi import FastAPI from probirka import Probirka, make_fastapi_endpoint app = FastAPI() - probirka_instance = Probirka() -fastapi_endpoint = make_fastapi_endpoint(probirka_instance) -app.add_api_route("/run", fastapi_endpoint) +# Define some health checks +@probirka_instance.add(name="api") +async def check_api(): + return True + + +# Create and add the endpoint +fastapi_endpoint = make_fastapi_endpoint(probirka_instance) +app.add_route("/health", fastapi_endpoint) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000) ``` -## Integration with aiohttp - -You can integrate Probirka with aiohttp as follows: +### aiohttp Integration ```python from aiohttp import web from probirka import Probirka, make_aiohttp_endpoint app = web.Application() - probirka_instance = Probirka() -aiohttp_endpoint = make_aiohttp_endpoint(probirka_instance) -app.router.add_get('/run', aiohttp_endpoint) +# Define some health checks +@probirka_instance.add(name="api") +async def check_api(): + return True + +# Create and add the endpoint +aiohttp_endpoint = make_aiohttp_endpoint(probirka_instance) +app.router.add_get('/health', aiohttp_endpoint) if __name__ == '__main__': web.run_app(app) ``` - -## Contributing - -Contributions are welcome! Please open an issue or submit a pull request on GitHub. - -## License - -This project is licensed under the MIT License. diff --git a/justfile b/justfile index 2e563b1..28f68c5 100644 --- a/justfile +++ b/justfile @@ -17,4 +17,4 @@ fix: uv run ruff check --fix --unsafe-fixes {{ SOURCE_PATH }} tests: - uv run pytest tests/ + uv run pytest --cov=probirka --cov-report lcov:tests.lcov tests/ diff --git a/probirka/_probes.py b/probirka/_probes.py index 91955b8..00d91ed 100644 --- a/probirka/_probes.py +++ b/probirka/_probes.py @@ -1,43 +1,76 @@ -from abc import ABC, abstractmethod +from abc import abstractmethod from asyncio import iscoroutinefunction, wait_for -from datetime import datetime -from typing import Callable, Optional, Protocol +from datetime import datetime, timedelta +from typing import Callable, Optional, Union, Dict, Any, Protocol from probirka._results import ProbeResult -class Probe( - Protocol, -): +class Probe(Protocol): """ - Protocol for a probe that can be run to check health. + Protocol defining the interface for a probe. """ - async def run_check( - self, - ) -> ProbeResult: ... + @property + def info(self) -> Optional[Dict[str, Any]]: ... + + def add_info(self, name: str, value: Any) -> None: ... + + async def run_check(self) -> ProbeResult: ... -class ProbeBase( - ABC, -): +class ProbeBase: """ - Abstract base class for a probe. + Base implementation of a probe. """ def __init__( self, name: Optional[str] = None, timeout: Optional[int] = None, + success_ttl: Optional[Union[int, timedelta]] = None, + failed_ttl: Optional[Union[int, timedelta]] = None, ) -> None: """ Initialize the probe. :param name: The name of the probe. :param timeout: The timeout for the probe. + :param success_ttl: Cache duration for successful results. If None, successful results are not cached. + :param failed_ttl: Cache duration for failed results. If None, failed results are not cached. """ self._timeout = timeout self._name = name or self.__class__.__name__ + self._success_ttl = timedelta(seconds=success_ttl) if isinstance(success_ttl, int) else success_ttl + self._failed_ttl = timedelta(seconds=failed_ttl) if isinstance(failed_ttl, int) else failed_ttl + self._last_result: Optional[ProbeResult] = None + self._info: Optional[Dict[str, Any]] = None + + def add_info( + self, + name: str, + value: Any, + ) -> None: + """ + Add information to the probe result. + + :param name: The name of the information. + :param value: The value of the information. + """ + if self._info is None: + self._info = {} + self._info[name] = value + + @property + def info( + self, + ) -> Optional[Dict[str, Any]]: + """ + Get the information added to the probe result. + + :return: The information added to the probe result. + """ + return self._info @abstractmethod async def _check( @@ -58,7 +91,25 @@ async def run_check( :return: The result of the check. """ - started_at = datetime.now() + now = datetime.now() + + if self._last_result is not None: + last_check_time = self._last_result.started_at + self._last_result.elapsed + ttl = self._success_ttl if self._last_result.ok else self._failed_ttl + if ttl is not None: + cache_until = last_check_time + ttl + if now < cache_until: + return ProbeResult( + ok=self._last_result.ok, + started_at=self._last_result.started_at, + elapsed=self._last_result.elapsed, + name=self._last_result.name, + error=self._last_result.error, + info=self._last_result.info or {}, + cached=bool(self._success_ttl is not None or self._failed_ttl is not None), + ) + + started_at = now error = None task = self._check() try: @@ -71,38 +122,44 @@ async def run_check( except Exception as exc: result = False error = str(exc) - finally: - task.close() - return ProbeResult( + + probe_result = ProbeResult( ok=False if result is None else result, started_at=started_at, elapsed=datetime.now() - started_at, name=self._name, error=error, + info=self._info or {}, + cached=False, ) + if self._success_ttl is not None or self._failed_ttl is not None: + ttl = self._success_ttl if probe_result.ok else self._failed_ttl + if ttl is not None: + self._last_result = probe_result + + return probe_result -class CallableProbe( - ProbeBase, -): + +class CallableProbe(ProbeBase): """ A probe that wraps a callable function. """ def __init__( - self, func: Callable[[], Optional[bool]], name: Optional[str] = None, timeout: Optional[int] = None + self, + func: Callable[[], Optional[bool]], + name: Optional[str] = None, + timeout: Optional[int] = None, + success_ttl: Optional[Union[int, timedelta]] = None, + failed_ttl: Optional[Union[int, timedelta]] = None, ) -> None: - """ - Initialize the callable probe. - - :param func: The callable function to wrap. - :param name: The name of the probe. - :param timeout: The timeout for the probe. - """ self._func = func super().__init__( name=name or func.__name__, timeout=timeout, + success_ttl=success_ttl, + failed_ttl=failed_ttl, ) async def _check( diff --git a/probirka/_probirka.py b/probirka/_probirka.py index 35eed4b..f08c717 100644 --- a/probirka/_probirka.py +++ b/probirka/_probirka.py @@ -1,6 +1,6 @@ from asyncio import gather, wait_for from collections import defaultdict -from datetime import datetime +from datetime import datetime, timedelta from typing import Any, Callable, Dict, List, Optional, Sequence, Union from probirka._probes import CallableProbe, Probe @@ -14,13 +14,20 @@ class Probirka: def __init__( self, + success_ttl: Optional[Union[int, timedelta]] = None, + failed_ttl: Optional[Union[int, timedelta]] = None, ) -> None: """ Initialize the Probirka instance. + + :param success_ttl: Default cache duration for successful results. If None, successful results are not cached. + :param failed_ttl: Default cache duration for failed results. If None, failed results are not cached. """ self._required_probes: List[Probe] = [] self._optional_probes: Dict[str, List[Probe]] = defaultdict(list) - self._info: Dict[str, Any] = {} + self._info: Optional[Dict[str, Any]] = None + self._success_ttl = timedelta(seconds=success_ttl) if isinstance(success_ttl, int) else success_ttl + self._failed_ttl = timedelta(seconds=failed_ttl) if isinstance(failed_ttl, int) else failed_ttl def add_info( self, @@ -33,16 +40,18 @@ def add_info( :param name: The name of the information. :param value: The value of the information. """ + if self._info is None: + self._info = {} self._info[name] = value @property def info( self, - ) -> Dict[str, Any]: + ) -> Optional[Dict[str, Any]]: """ Get the information added to the health check result. - :return: A dictionary containing the information. + :return: The information added to the health check result. """ return self._info @@ -54,8 +63,8 @@ def add_probes( """ Add probes to the health check. - :param probes: The probes to add. - :param groups: The groups to which the probes belong. + :param probes: Probes to add + :param groups: Groups for optional probes. Probes without groups are required. """ if groups: if isinstance(groups, str): @@ -70,14 +79,18 @@ def add( name: Optional[str] = None, timeout: Optional[int] = None, groups: Union[str, List[str]] = '', + success_ttl: Optional[Union[int, timedelta]] = None, + failed_ttl: Optional[Union[int, timedelta]] = None, ) -> Callable: """ Decorator to add a callable as a probe. - :param name: The name of the probe. - :param timeout: The timeout for the probe. - :param groups: The groups to which the probe belongs. - :return: The decorated function. + :param name: Probe name + :param timeout: Probe timeout in seconds + :param groups: Groups for optional probes. Probes without groups are required. + :param success_ttl: Cache duration for successful results. If None, uses the global success_ttl setting. + :param failed_ttl: Cache duration for failed results. If None, uses the global failed_ttl setting. + :return: Decorated function """ def _wrapper(func: Callable) -> Any: @@ -86,6 +99,8 @@ def _wrapper(func: Callable) -> Any: func=func, name=name, timeout=timeout, + success_ttl=success_ttl or self._success_ttl, + failed_ttl=failed_ttl or self._failed_ttl, ), groups=groups, ) @@ -99,11 +114,11 @@ async def _inner_run( skip_required: bool, ) -> Sequence[ProbeResult]: """ - Run the probes and gather the results. + Run probes and gather results. - :param with_groups: The groups of probes to run. - :param skip_required: Whether to skip the required probes. - :return: A sequence of probe results. + :param with_groups: Groups to run. Required probes run unless skip_required=True + :param skip_required: Skip probes without groups + :return: Sequence of probe results """ tasks = [] if skip_required else [probe.run_check() for probe in self._required_probes] for group in with_groups: @@ -120,12 +135,12 @@ async def run( skip_required: bool = False, ) -> HealthCheckResult: """ - Run the health check and return the result. + Run health check and return results. - :param timeout: The timeout for the health check. - :param with_groups: The groups of probes to run. - :param skip_required: Whether to skip the required probes. - :return: The health check result. + :param timeout: Overall timeout in seconds + :param with_groups: Groups to run. Required probes run unless skip_required=True + :param skip_required: Skip probes without groups + :return: Health check result """ if with_groups and isinstance(with_groups, str): with_groups = [with_groups] diff --git a/probirka/_results.py b/probirka/_results.py index 0ef8af7..54e6f1f 100644 --- a/probirka/_results.py +++ b/probirka/_results.py @@ -10,12 +10,14 @@ class ProbeResult: elapsed: timedelta name: str error: Optional[str] + info: Optional[Dict[str, Any]] + cached: Optional[bool] @dataclass(frozen=True) class HealthCheckResult: ok: bool - info: Dict[str, Any] + info: Optional[Dict[str, Any]] started_at: datetime total_elapsed: timedelta checks: Sequence[ProbeResult] diff --git a/pyproject.toml b/pyproject.toml index 2fdf08b..b3c2293 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ dev = [ "aiohttp>=3.10.11", "bandit>=1.7.10", "fastapi>=0.115.6", + "httpx>=0.28.1", "mypy>=1.13.0", "pre-commit>=3.5.0", "pytest>=8.3.4", diff --git a/tests/conftest.py b/tests/conftest.py index 61bc56c..450db64 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,7 @@ +import pytest + +pytest_plugins = ('pytest_asyncio',) + from typing import Callable, Optional, Union from unittest.mock import MagicMock diff --git a/tests/test_aiohttp.py b/tests/test_aiohttp.py new file mode 100644 index 0000000..ee60ca8 --- /dev/null +++ b/tests/test_aiohttp.py @@ -0,0 +1,144 @@ +import json +from datetime import datetime, timedelta +from typing import AsyncGenerator + +import pytest +from aiohttp import web +from aiohttp.test_utils import TestClient, TestServer + +from probirka import Probirka, ProbeBase +from probirka._aiohttp import make_aiohttp_endpoint +from probirka._results import HealthCheckResult, ProbeResult + + +class SuccessProbe(ProbeBase): + async def _check(self) -> bool: + return True + + +class FailureProbe(ProbeBase): + async def _check(self) -> bool: + return False + + +@pytest.fixture +def probirka() -> Probirka: + return Probirka() + + +@pytest.mark.asyncio +async def test_successful_response(probirka: Probirka) -> None: + # Подготовка + app = web.Application() + endpoint = make_aiohttp_endpoint(probirka) + app.router.add_get("/health", endpoint) + server = TestServer(app) + + probirka.add_info("some_field", "value") + probirka.add_probes(SuccessProbe()) + + async with TestClient(server) as client: + # Выполнение + async with client.get("/health") as response: + # Проверка + assert response.status == 200 + response_data = await response.json() + assert response_data["ok"] is True + assert response_data["info"]["some_field"] == "value" + assert len(response_data["checks"]) == 1 + assert response_data["checks"][0]["ok"] is True + + +@pytest.mark.asyncio +async def test_error_response(probirka: Probirka) -> None: + # Подготовка + app = web.Application() + endpoint = make_aiohttp_endpoint(probirka) + app.router.add_get("/health", endpoint) + server = TestServer(app) + + probirka.add_probes(FailureProbe()) + + async with TestClient(server) as client: + # Выполнение + async with client.get("/health") as response: + # Проверка + assert response.status == 500 + response_data = await response.json() + assert response_data["ok"] is False + assert len(response_data["checks"]) == 1 + assert response_data["checks"][0]["ok"] is False + + +@pytest.mark.asyncio +async def test_custom_status_codes(probirka: Probirka) -> None: + # Подготовка + app = web.Application() + endpoint = make_aiohttp_endpoint( + probirka, + success_code=201, + error_code=400 + ) + app.router.add_get("/health", endpoint) + server = TestServer(app) + + probirka.add_probes(SuccessProbe()) + + async with TestClient(server) as client: + # Выполнение + async with client.get("/health") as response: + # Проверка + assert response.status == 201 + response_data = await response.json() + assert response_data["ok"] is True + + +@pytest.mark.asyncio +async def test_without_results(probirka: Probirka) -> None: + # Подготовка + app = web.Application() + endpoint = make_aiohttp_endpoint( + probirka, + return_results=False + ) + app.router.add_get("/health", endpoint) + server = TestServer(app) + + probirka.add_probes(SuccessProbe()) + + async with TestClient(server) as client: + # Выполнение + async with client.get("/health") as response: + # Проверка + assert response.status == 200 + assert await response.text() == "" + + +@pytest.mark.asyncio +async def test_with_custom_parameters(probirka: Probirka) -> None: + # Подготовка + success_probe_1 = SuccessProbe() + success_probe_2 = SuccessProbe() + + probirka.add_probes(success_probe_1) # required probe + probirka.add_probes(success_probe_2, groups=["group1"]) # optional probe + + app = web.Application() + endpoint = make_aiohttp_endpoint( + probirka, + timeout=30, + with_groups=["group1"], + skip_required=True + ) + app.router.add_get("/health", endpoint) + server = TestServer(app) + + async with TestClient(server) as client: + # Выполнение + async with client.get("/health") as response: + # Проверка + assert response.status == 200 + response_data = await response.json() + assert response_data["ok"] is True + # Проверяем, что запустился только один проб из группы group1 + assert len(response_data["checks"]) == 1 diff --git a/tests/test_fastapi.py b/tests/test_fastapi.py new file mode 100644 index 0000000..99580be --- /dev/null +++ b/tests/test_fastapi.py @@ -0,0 +1,132 @@ +import pytest +from fastapi import FastAPI, status +from fastapi.testclient import TestClient + +from probirka import Probirka, ProbeBase +from probirka._fastapi import make_fastapi_endpoint + + +class SuccessProbe(ProbeBase): + async def _check(self) -> bool: + return True + + +class FailureProbe(ProbeBase): + async def _check(self) -> bool: + return False + + +@pytest.fixture +def probirka() -> Probirka: + return Probirka() + + +@pytest.fixture +def test_client(probirka: Probirka) -> TestClient: + app = FastAPI() + endpoint = make_fastapi_endpoint(probirka) + app.add_api_route("/health", endpoint) + return TestClient(app) + + +def test_successful_response(test_client: TestClient, probirka: Probirka) -> None: + # Подготовка + probirka.add_info("some_field", "value") + probirka.add_probes(SuccessProbe()) + + # Выполнение + response = test_client.get("/health") + + # Проверка + assert response.status_code == status.HTTP_200_OK + response_data = response.json() + assert response_data["ok"] is True + assert response_data["info"]["some_field"] == "value" + assert len(response_data["checks"]) == 1 + assert response_data["checks"][0]["ok"] is True + + +def test_error_response(test_client: TestClient, probirka: Probirka) -> None: + # Подготовка + probirka.add_probes(FailureProbe()) + + # Выполнение + response = test_client.get("/health") + + # Проверка + assert response.status_code == status.HTTP_500_INTERNAL_SERVER_ERROR + response_data = response.json() + assert response_data["ok"] is False + assert len(response_data["checks"]) == 1 + assert response_data["checks"][0]["ok"] is False + + +def test_custom_status_codes(probirka: Probirka) -> None: + # Подготовка + app = FastAPI() + endpoint = make_fastapi_endpoint( + probirka, + success_code=status.HTTP_201_CREATED, + error_code=status.HTTP_400_BAD_REQUEST + ) + app.add_api_route("/health", endpoint) + client = TestClient(app) + + probirka.add_probes(SuccessProbe()) + + # Выполнение + response = client.get("/health") + + # Проверка + assert response.status_code == status.HTTP_201_CREATED + response_data = response.json() + assert response_data["ok"] is True + + +def test_without_results(probirka: Probirka) -> None: + # Подготовка + app = FastAPI() + endpoint = make_fastapi_endpoint( + probirka, + return_results=False + ) + app.add_api_route("/health", endpoint) + client = TestClient(app) + + probirka.add_probes(SuccessProbe()) + + # Выполнение + response = client.get("/health") + + # Проверка + assert response.status_code == status.HTTP_200_OK + assert response.content == b"" + + +def test_with_custom_parameters(probirka: Probirka) -> None: + # Подготовка + success_probe_1 = SuccessProbe() + success_probe_2 = SuccessProbe() + + probirka.add_probes(success_probe_1) # required probe + probirka.add_probes(success_probe_2, groups=["group1"]) # optional probe + + app = FastAPI() + endpoint = make_fastapi_endpoint( + probirka, + timeout=30, + with_groups=["group1"], + skip_required=True + ) + app.add_api_route("/health", endpoint) + client = TestClient(app) + + # Выполнение + response = client.get("/health") + + # Проверка + assert response.status_code == status.HTTP_200_OK + response_data = response.json() + assert response_data["ok"] is True + # Проверяем, что запустился только один проб из группы group1 + assert len(response_data["checks"]) == 1 diff --git a/tests/test_probes.py b/tests/test_probes.py index 3a88202..30c8016 100644 --- a/tests/test_probes.py +++ b/tests/test_probes.py @@ -1,9 +1,14 @@ -from typing import Callable, Optional, Union +from typing import Callable, Optional, Union, Any from unittest.mock import MagicMock import pytest +import asyncio +from datetime import datetime, timedelta +from unittest.mock import AsyncMock, patch from probirka import Probe +from probirka._probes import ProbeBase, CallableProbe +from probirka._results import ProbeResult @pytest.mark.parametrize( @@ -24,3 +29,220 @@ async def test_run_check( probe = make_testing_probe(probe_result) results = await probe.run_check() assert results.ok == is_ok, results + + +class TestProbeBase: + class ConcreteProbe(ProbeBase): + async def _check(self) -> bool: + return True + + def test_init_with_default_values(self) -> None: + probe = self.ConcreteProbe() + assert probe._name == "ConcreteProbe" + assert probe._timeout is None + + def test_init_with_custom_values(self) -> None: + probe = self.ConcreteProbe(name="CustomProbe", timeout=5) + assert probe._name == "CustomProbe" + assert probe._timeout == 5 + + @pytest.mark.asyncio + async def test_run_check_success(self) -> None: + probe = self.ConcreteProbe() + result = await probe.run_check() + + assert isinstance(result, ProbeResult) + assert result.ok is True + assert isinstance(result.started_at, datetime) + assert isinstance(result.elapsed, timedelta) + assert result.name == "ConcreteProbe" + assert result.error is None + + @pytest.mark.asyncio + async def test_run_check_with_timeout(self) -> None: + class SlowProbe(ProbeBase): + async def _check(self) -> bool: + await asyncio.sleep(2) + return True + + probe = SlowProbe(timeout=1) + result = await probe.run_check() + + assert result.ok is False + assert result.error is not None + + @pytest.mark.asyncio + async def test_run_check_with_exception(self) -> None: + class FailingProbe(ProbeBase): + async def _check(self) -> bool: + raise ValueError("Test error") + + probe = FailingProbe() + result = await probe.run_check() + + assert result.ok is False + assert result.error == "Test error" + + +class TestCallableProbe: + def test_init_with_sync_function(self) -> None: + def test_func() -> bool: + return True + + probe = CallableProbe(test_func) + assert probe._name == "test_func" + assert probe._func == test_func + + def test_init_with_async_function(self) -> None: + async def test_func() -> bool: + return True + + probe = CallableProbe(test_func) + assert probe._name == "test_func" + assert probe._func == test_func + + @pytest.mark.asyncio + async def test_check_with_sync_function(self) -> None: + def test_func() -> bool: + return True + + probe = CallableProbe(test_func) + result = await probe._check() + assert result is True + + @pytest.mark.asyncio + async def test_check_with_async_function(self) -> None: + async def test_func() -> bool: + return True + + probe = CallableProbe(test_func) + result = await probe._check() + assert result is True + + @pytest.mark.asyncio + async def test_run_check_with_sync_function(self) -> None: + def test_func() -> bool: + return True + + probe = CallableProbe(test_func) + result = await probe.run_check() + + assert isinstance(result, ProbeResult) + assert result.ok is True + assert result.name == "test_func" + + @pytest.mark.asyncio + async def test_run_check_with_async_function(self) -> None: + async def test_func() -> bool: + return True + + probe = CallableProbe(test_func) + result = await probe.run_check() + + assert isinstance(result, ProbeResult) + assert result.ok is True + assert result.name == "test_func" + + +@pytest.mark.asyncio +async def test_probe_caching() -> None: + class TestProbe(ProbeBase): + def __init__(self, success_ttl: Optional[int] = None, failed_ttl: Optional[int] = None): + super().__init__(success_ttl=success_ttl, failed_ttl=failed_ttl) + self._counter = 0 + + async def _check(self) -> bool: + self._counter += 1 + return True + + # Тест без кэширования + probe = TestProbe() + result1 = await probe.run_check() + result2 = await probe.run_check() + assert result1.ok is True + assert result2.ok is True + assert result1.cached is False + assert result2.cached is False + assert probe._counter == 2 + + # Тест с кэшированием успешного результата + probe = TestProbe(success_ttl=1) + result1 = await probe.run_check() + result2 = await probe.run_check() + assert result1.ok is True + assert result2.ok is True + assert result1.cached is False + assert result2.cached is True + assert probe._counter == 1 + + # Тест с кэшированием неуспешного результата + class FailingProbe(ProbeBase): + def __init__(self, failed_ttl: Optional[int] = None): + super().__init__(failed_ttl=failed_ttl) + self._counter = 0 + + async def _check(self) -> bool: + self._counter += 1 + return False + + probe = FailingProbe(failed_ttl=1) + result1 = await probe.run_check() + result2 = await probe.run_check() + assert result1.ok is False + assert result2.ok is False + assert result1.cached is False + assert result2.cached is True + assert probe._counter == 1 + + +@pytest.mark.asyncio +async def test_probe_info() -> None: + class TestProbe(ProbeBase): + async def _check(self) -> bool: + self.add_info("test_key", "test_value") + return True + + probe = TestProbe() + result = await probe.run_check() + + assert result.ok is True + assert result.info == {"test_key": "test_value"} + + +@pytest.mark.asyncio +async def test_probe_info_caching() -> None: + class TestProbe(ProbeBase): + def __init__(self, success_ttl: Optional[int] = None): + super().__init__(success_ttl=success_ttl) + self._counter = 0 + + async def _check(self) -> bool: + self._counter += 1 + self.add_info("counter", self._counter) + return True + + # Тест с кэшированием + probe = TestProbe(success_ttl=1) + result1 = await probe.run_check() + result2 = await probe.run_check() + + assert result1.ok is True + assert result2.ok is True + assert result1.cached is False + assert result2.cached is True + assert result1.info == {"counter": 1} + assert result2.info == {"counter": 1} + assert probe._counter == 1 + + +@pytest.mark.asyncio +async def test_probe_info_empty() -> None: + class TestProbe(ProbeBase): + async def _check(self) -> bool: + return True + + probe = TestProbe() + result = await probe.run_check() + + assert result.ok is True + assert result.info == {} diff --git a/tests/test_probirka.py b/tests/test_probirka.py index aba7885..cc18862 100644 --- a/tests/test_probirka.py +++ b/tests/test_probirka.py @@ -28,9 +28,54 @@ def _check_1() -> bool: def _check_2() -> bool: return False - results = await checks.run(with_groups='optional') - assert len(results.checks) == 2 # noqa + # Проверяем, что при запуске только опциональных проверок запускается только одна проверка + results = await checks.run(with_groups='optional', skip_required=True) + assert len(results.checks) == 1 + assert results.checks[0].ok is False + # Проверяем, что при запуске всех проверок запускается только обязательная проверка results = await checks.run() assert len(results.checks) == 1 assert results.checks[0].ok is True + + +@pytest.mark.asyncio +async def test_probirka_caching() -> None: + checks = Probirka() + counter = 0 + + @checks.add(success_ttl=1) + def _check_1() -> bool: + nonlocal counter + counter += 1 + return True + + # Первый запуск + results = await checks.run() + assert results.checks[0].ok is True + assert counter == 1 + + # Второй запуск (должен использовать кэш) + results = await checks.run() + assert results.checks[0].ok is True + assert counter == 1 + + # Проверка глобальных настроек кэширования + checks2 = Probirka(success_ttl=1, failed_ttl=1) + counter2 = 0 + + @checks2.add() + def _check_2() -> bool: + nonlocal counter2 + counter2 += 1 + return True + + # Первый запуск + results = await checks2.run() + assert results.checks[0].ok is True + assert counter2 == 1 + + # Второй запуск (должен использовать кэш) + results = await checks2.run() + assert results.checks[0].ok is True + assert counter2 == 1 diff --git a/uv.lock b/uv.lock index ae451d6..5968b46 100644 --- a/uv.lock +++ b/uv.lock @@ -1,4 +1,5 @@ version = 1 +revision = 1 requires-python = ">=3.8" resolution-markers = [ "python_full_version < '3.13'", @@ -195,7 +196,7 @@ name = "bandit" version = "1.7.10" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "platform_system == 'Windows'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, { name = "pyyaml" }, { name = "rich" }, { name = "stevedore" }, @@ -404,7 +405,7 @@ name = "click" version = "8.1.7" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "platform_system == 'Windows'" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/96/d3/f04c7bfcf5c1862a2a5b845c6b2b360488cf47af55dfa79c98f6a6bf98b5/click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de", size = 336121 } wheels = [ @@ -519,7 +520,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/98/65/13d9e76ca19b0ba5603d71ac8424b5694415b348e719db277b5edc985ff5/cryptography-44.0.0-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:761817a3377ef15ac23cd7834715081791d4ec77f9297ee694ca1ee9c2c7e5eb", size = 3915420 }, { url = "https://files.pythonhosted.org/packages/b1/07/40fe09ce96b91fc9276a9ad272832ead0fddedcba87f1190372af8e3039c/cryptography-44.0.0-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3c672a53c0fb4725a29c303be906d3c1fa99c32f58abe008a82705f9ee96f40b", size = 4154498 }, { url = "https://files.pythonhosted.org/packages/75/ea/af65619c800ec0a7e4034207aec543acdf248d9bffba0533342d1bd435e1/cryptography-44.0.0-cp37-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:4ac4c9f37eba52cb6fbeaf5b59c152ea976726b865bd4cf87883a7e7006cc543", size = 3932569 }, - { url = "https://files.pythonhosted.org/packages/4e/d5/9cc182bf24c86f542129565976c21301d4ac397e74bf5a16e48241aab8a6/cryptography-44.0.0-cp37-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:60eb32934076fa07e4316b7b2742fa52cbb190b42c2df2863dbc4230a0a9b385", size = 4164756 }, { url = "https://files.pythonhosted.org/packages/c7/af/d1deb0c04d59612e3d5e54203159e284d3e7a6921e565bb0eeb6269bdd8a/cryptography-44.0.0-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ed3534eb1090483c96178fcb0f8893719d96d5274dfde98aa6add34614e97c8e", size = 4016721 }, { url = "https://files.pythonhosted.org/packages/bd/69/7ca326c55698d0688db867795134bdfac87136b80ef373aaa42b225d6dd5/cryptography-44.0.0-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:f3f6fdfa89ee2d9d496e2c087cebef9d4fcbb0ad63c40e821b39f74bf48d9c5e", size = 4240915 }, { url = "https://files.pythonhosted.org/packages/ef/d4/cae11bf68c0f981e0413906c6dd03ae7fa864347ed5fac40021df1ef467c/cryptography-44.0.0-cp37-abi3-win32.whl", hash = "sha256:eb33480f1bad5b78233b0ad3e1b0be21e8ef1da745d8d2aecbb20671658b9053", size = 2757925 }, @@ -530,7 +530,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d0/c7/c656eb08fd22255d21bc3129625ed9cd5ee305f33752ef2278711b3fa98b/cryptography-44.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:c5eb858beed7835e5ad1faba59e865109f3e52b3783b9ac21e7e47dc5554e289", size = 3915417 }, { url = "https://files.pythonhosted.org/packages/ef/82/72403624f197af0db6bac4e58153bc9ac0e6020e57234115db9596eee85d/cryptography-44.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:f53c2c87e0fb4b0c00fa9571082a057e37690a8f12233306161c8f4b819960b7", size = 4155160 }, { url = "https://files.pythonhosted.org/packages/a2/cd/2f3c440913d4329ade49b146d74f2e9766422e1732613f57097fea61f344/cryptography-44.0.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:9e6fc8a08e116fb7c7dd1f040074c9d7b51d74a8ea40d4df2fc7aa08b76b9e6c", size = 3932331 }, - { url = "https://files.pythonhosted.org/packages/31/d9/90409720277f88eb3ab72f9a32bfa54acdd97e94225df699e7713e850bd4/cryptography-44.0.0-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:9abcc2e083cbe8dde89124a47e5e53ec38751f0d7dfd36801008f316a127d7ba", size = 4165207 }, { url = "https://files.pythonhosted.org/packages/7f/df/8be88797f0a1cca6e255189a57bb49237402b1880d6e8721690c5603ac23/cryptography-44.0.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:d2436114e46b36d00f8b72ff57e598978b37399d2786fd39793c36c6d5cb1c64", size = 4017372 }, { url = "https://files.pythonhosted.org/packages/af/36/5ccc376f025a834e72b8e52e18746b927f34e4520487098e283a719c205e/cryptography-44.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a01956ddfa0a6790d594f5b34fc1bfa6098aca434696a03cfdbe469b8ed79285", size = 4239657 }, { url = "https://files.pythonhosted.org/packages/46/b0/f4f7d0d0bcfbc8dd6296c1449be326d04217c57afb8b2594f017eed95533/cryptography-44.0.0-cp39-abi3-win32.whl", hash = "sha256:eca27345e1214d1b9f9490d200f9db5a874479be914199194e746c893788d417", size = 2758672 }, @@ -696,6 +695,43 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c6/c8/a5be5b7550c10858fcf9b0ea054baccab474da77d37f1e828ce043a3a5d4/frozenlist-1.5.0-py3-none-any.whl", hash = "sha256:d994863bba198a4a518b467bb971c56e1db3f180a25c6cf7bb1949c267f748c3", size = 11901 }, ] +[[package]] +name = "h11" +version = "0.14.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f5/38/3af3d3633a34a3316095b39c8e8fb4853a28a536e55d347bd8d8e9a14b03/h11-0.14.0.tar.gz", hash = "sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d", size = 100418 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/95/04/ff642e65ad6b90db43e668d70ffb6736436c7ce41fcc549f4e9472234127/h11-0.14.0-py3-none-any.whl", hash = "sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761", size = 58259 }, +] + +[[package]] +name = "httpcore" +version = "1.0.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "h11" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6a/41/d7d0a89eb493922c37d343b607bc1b5da7f5be7e383740b4753ad8943e90/httpcore-1.0.7.tar.gz", hash = "sha256:8551cb62a169ec7162ac7be8d4817d561f60e08eaa485234898414bb5a8a0b4c", size = 85196 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/87/f5/72347bc88306acb359581ac4d52f23c0ef445b57157adedb9aee0cd689d2/httpcore-1.0.7-py3-none-any.whl", hash = "sha256:a3fff8f43dc260d5bd363d9f9cf1830fa3a458b332856f34282de498ed420edd", size = 78551 }, +] + +[[package]] +name = "httpx" +version = "0.28.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "anyio" }, + { name = "certifi" }, + { name = "httpcore" }, + { name = "idna" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517 }, +] + [[package]] name = "identify" version = "2.6.1" @@ -1052,6 +1088,7 @@ dev = [ { name = "aiohttp" }, { name = "bandit" }, { name = "fastapi" }, + { name = "httpx" }, { name = "mypy" }, { name = "pre-commit" }, { name = "pytest" }, @@ -1069,6 +1106,7 @@ dev = [ { name = "aiohttp", specifier = ">=3.10.11" }, { name = "bandit", specifier = ">=1.7.10" }, { name = "fastapi", specifier = ">=0.115.6" }, + { name = "httpx", specifier = ">=0.28.1" }, { name = "mypy", specifier = ">=1.13.0" }, { name = "pre-commit", specifier = ">=3.5.0" }, { name = "pytest", specifier = ">=8.3.4" },