diff options
30 files changed, 795 insertions, 251 deletions
diff --git a/.github/workflows/build-on-release.yml b/.github/workflows/build-on-release.yml new file mode 100644 index 0000000..498393f --- /dev/null +++ b/.github/workflows/build-on-release.yml @@ -0,0 +1,200 @@ +name: Build and Release + +on: + release: + types: [created] + +jobs: + build: + strategy: + matrix: + os: [ubuntu-latest, windows-latest, macos-latest] + + runs-on: ${{ matrix.os }} + + steps: + - uses: actions/checkout@v2 + + - name: Set up Python + uses: actions/setup-python@v2 + with: + python-version: '3.9' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install pyinstaller + + - name: Install dependancies (Linux) + if: matrix.os == 'ubuntu-latest' + run: | + sudo apt-get install -y binutils + + - name: Install the Apple certificate and provisioning profile + if: matrix.os == 'macos-latest' + env: + BUILD_CERTIFICATE_BASE64: ${{ secrets.BUILD_CERTIFICATE_BASE64 }} + P12_PASSWORD: ${{ secrets.P12_PASSWORD }} + BUILD_PROVISION_PROFILE_BASE64: ${{ secrets.BUILD_PROVISION_PROFILE_BASE64 }} + KEYCHAIN_PASSWORD: ${{ secrets.KEYCHAIN_PASSWORD }} + run: | + # create variables + CERTIFICATE_PATH=$RUNNER_TEMP/build_certificate.p12 + PP_PATH=$RUNNER_TEMP/build_pp.mobileprovision + KEYCHAIN_PATH=$RUNNER_TEMP/app-signing.keychain-db + + # import certificate and provisioning profile from secrets + echo -n "$BUILD_CERTIFICATE_BASE64" | base64 --decode -o $CERTIFICATE_PATH + echo -n "$BUILD_PROVISION_PROFILE_BASE64" | base64 --decode -o $PP_PATH + + # create temporary keychain + security create-keychain -p "$KEYCHAIN_PASSWORD" $KEYCHAIN_PATH + security set-keychain-settings -lut 21600 $KEYCHAIN_PATH + security unlock-keychain -p "$KEYCHAIN_PASSWORD" $KEYCHAIN_PATH + + # import certificate to keychain + security import $CERTIFICATE_PATH -P "$P12_PASSWORD" -A -t cert -f pkcs12 -k $KEYCHAIN_PATH + security list-keychain -d user -s $KEYCHAIN_PATH + + # apply provisioning profile + mkdir -p ~/Library/MobileDevice/Provisioning\ Profiles + cp $PP_PATH ~/Library/MobileDevice/Provisioning\ Profiles + + - name: Set version in code (Unix) + if: matrix.os == 'macos-latest' || matrix.os == 'ubuntu-latest' + run: | + awk 'NR==3{$0="__version__ = \"'${{ github.ref_name }}'\""}1' ./robusta_krr/__init__.py > temp && mv temp ./robusta_krr/__init__.py + cat ./robusta_krr/__init__.py + + - name: Set version in code (Windows) + if: matrix.os == 'windows-latest' + run: | + $content = Get-Content -Path .\robusta_krr\__init__.py + $content[2] = "__version__=`"$($env:GITHUB_REF_NAME)`"" + $content | Out-File -FilePath .\robusta_krr\__init__.py -Encoding ascii + Get-Content .\robusta_krr\__init__.py + shell: pwsh + env: + GITHUB_REF_NAME: ${{ github.ref_name }} + + + - name: Build with PyInstaller + shell: bash + run: | + pyinstaller krr.py + mkdir -p ./dist/krr/grapheme/data + cp $(python -c "import grapheme; print(grapheme.__path__[0] + '/data/grapheme_break_property.json')") ./dist/krr/grapheme/data/grapheme_break_property.json + + - name: Zip the application (Unix) + if: matrix.os == 'macos-latest' || matrix.os == 'ubuntu-latest' + run: | + cd dist + zip -r krr-${{ matrix.os }}-${{ github.ref_name }}.zip krr + mv krr-${{ matrix.os }}-${{ github.ref_name }}.zip ../ + cd .. + + - name: Zip the application (Windows) + if: matrix.os == 'windows-latest' + run: | + Set-Location -Path dist + Compress-Archive -Path krr -DestinationPath krr-${{ matrix.os }}-${{ github.ref_name }}.zip -Force + Move-Item -Path krr-${{ matrix.os }}-${{ github.ref_name }}.zip -Destination ..\ + Set-Location -Path .. + + - name: Upload Release Asset + uses: actions/upload-release-asset@v1.0.2 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + upload_url: ${{ github.event.release.upload_url }} + asset_path: ./krr-${{ matrix.os }}-${{ github.ref_name }}.zip + asset_name: krr-${{ matrix.os }}-${{ github.ref_name }}.zip + asset_content_type: application/octet-stream + + - name: Upload build as artifact + uses: actions/upload-artifact@v2 + with: + name: krr-${{ matrix.os }}-${{ github.ref_name }} + path: ./krr-${{ matrix.os }}-${{ github.ref_name }}.zip + + - name: Clean up keychain and provisioning profile + if: (matrix.os == 'macos-latest') && always() + run: | + security delete-keychain $RUNNER_TEMP/app-signing.keychain-db + rm ~/Library/MobileDevice/Provisioning\ Profiles/build_pp.mobileprovision + + check-latest: + needs: build + runs-on: ubuntu-latest + outputs: + IS_LATEST: ${{ steps.check-latest.outputs.release == github.ref_name }} + steps: + - id: check-latest + uses: pozetroninc/github-action-get-latest-release@v0.7.0 + with: + token: ${{ secrets.GITHUB_TOKEN }} + repository: ${{ github.repository }} + excludes: prerelease, draft + + # Define MacOS hash job + mac-hash: + needs: check-latest + runs-on: ubuntu-latest + if: needs.check-latest.outputs.IS_LATEST + outputs: + MAC_BUILD_HASH: ${{ steps.calc-hash.outputs.MAC_BUILD_HASH }} + steps: + - name: Checkout Repository + uses: actions/checkout@v2 + - name: Download MacOS artifact + uses: actions/download-artifact@v2 + with: + name: krr-macos-latest-${{ github.ref_name }} + - name: Calculate hash + id: calc-hash + run: echo "::set-output name=MAC_BUILD_HASH::$(sha256sum krr-macos-latest-${{ github.ref_name }}.zip | awk '{print $1}')" + + # Define Linux hash job + linux-hash: + needs: check-latest + runs-on: ubuntu-latest + if: needs.check-latest.outputs.IS_LATEST + outputs: + LINUX_BUILD_HASH: ${{ steps.calc-hash.outputs.LINUX_BUILD_HASH }} + steps: + - name: Checkout Repository + uses: actions/checkout@v2 + - name: Download Linux artifact + uses: actions/download-artifact@v2 + with: + name: krr-ubuntu-latest-${{ github.ref_name }} + - name: Calculate hash + id: calc-hash + run: echo "::set-output name=LINUX_BUILD_HASH::$(sha256sum krr-ubuntu-latest-${{ github.ref_name }}.zip | awk '{print $1}')" + + # Define job to update homebrew formula + update-formula: + needs: [mac-hash, linux-hash] + runs-on: ubuntu-latest + steps: + - name: Checkout homebrew-krr repository + uses: actions/checkout@v2 + with: + repository: robusta-dev/homebrew-krr + token: ${{ secrets.MULTIREPO_GITHUB_TOKEN }} + - name: Update krr.rb formula + run: | + MAC_BUILD_HASH=${{ needs.mac-hash.outputs.MAC_BUILD_HASH }} + LINUX_BUILD_HASH=${{ needs.linux-hash.outputs.LINUX_BUILD_HASH }} + TAG_NAME=${{ github.ref_name }} + awk 'NR==6{$0=" url \"https://github.com/robusta-dev/krr/releases/download/'"$TAG_NAME"'/krr-macos-latest-'"$TAG_NAME"'.zip\""}1' ./Formula/krr.rb > temp && mv temp ./Formula/krr.rb + awk 'NR==7{$0=" sha256 \"'$MAC_BUILD_HASH'\""}1' ./Formula/krr.rb > temp && mv temp ./Formula/krr.rb + awk 'NR==9{$0=" url \"https://github.com/robusta-dev/krr/releases/download/'"$TAG_NAME"'/krr-linux-latest-'"$TAG_NAME"'.zip\""}1' ./Formula/krr.rb > temp && mv temp ./Formula/krr.rb + awk 'NR==10{$0=" sha256 \"'$LINUX_BUILD_HASH'\""}1' ./Formula/krr.rb > temp && mv temp ./Formula/krr.rb + - name: Commit and push changes + run: | + git config --local user.email "action@github.com" + git config --local user.name "GitHub Action" + git commit -am "Update formula for release ${TAG_NAME}" + git push diff --git a/.github/workflows/docker-build-on-tag.yml b/.github/workflows/docker-build-on-tag.yml new file mode 100644 index 0000000..97829f5 --- /dev/null +++ b/.github/workflows/docker-build-on-tag.yml @@ -0,0 +1,43 @@ +name: Docker Build and Push + +on: + push: + tags: + - '*' + +jobs: + build: + + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v2 + + - name: Set up gcloud CLI + uses: google-github-actions/setup-gcloud@v0.2.0 + with: + service_account_key: ${{ secrets.GCP_SA_KEY }} + project_id: genuine-flight-317411 + export_default_credentials: true + + # Configure Docker to use the gcloud command-line tool as a credential helper for authentication + - name: Configure Docker + run: |- + gcloud auth configure-docker us-central1-docker.pkg.dev + + - name: Verify gcloud configuration + run: |- + gcloud config get-value project + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v1 + + - name: Build and push Docker images + uses: docker/build-push-action@v2 + with: + context: . + platforms: linux/arm64,linux/amd64 + push: true + tags: us-central1-docker.pkg.dev/genuine-flight-317411/devel/krr:${{ github.ref_name }} + build-args: | + BUILDKIT_INLINE_CACHE=1 diff --git a/.github/workflows/pytest-on-push.yml b/.github/workflows/pytest-on-push.yml new file mode 100644 index 0000000..a867168 --- /dev/null +++ b/.github/workflows/pytest-on-push.yml @@ -0,0 +1,27 @@ +name: Pytest + +on: [push] + +jobs: + build: + + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v2 + + - name: Set up Python + uses: actions/setup-python@v2 + with: + python-version: '3.9' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -r requirements.txt + pip install -e . + pip install pytest + + - name: Test with pytest + run: | + pytest @@ -1,4 +1,5 @@ # Byte-compiled / optimized / DLL files +.idea/ __pycache__/ *.py[cod] *$py.class @@ -130,4 +131,4 @@ dmypy.json .DS_Store -robusta_lib
\ No newline at end of file +robusta_lib @@ -133,6 +133,40 @@ More features (like seeing graphs, based on which recommendations were made) com <p align="right">(<a href="#readme-top">back to top</a>)</p> +<!-- ADVANCED USAGE EXAMPLES --> + +### Slack integration + +Put cost savings on autopilot. Get notified in Slack about recommendations above X%. Send a weekly global report, or one report per team. + +![Slack Screen Shot][slack-screenshot] + +#### Prerequisites +* A Slack workspace + +#### Setup +1. [Install Robusta with Helm to your cluster and configure slack](https://docs.robusta.dev/master/installation.html) +2. Create your KRR slack playbook by adding the following to `generated_values.yaml`: +``` +customPlaybooks: +# Runs a weekly krr scan on the namespace devs-namespace and sends it to the configured slack channel +customPlaybooks: +- triggers: + - on_schedule: + fixed_delay_repeat: + repeat: 1 # number of times to run or -1 to run forever + seconds_delay: 604800 # 1 week + actions: + - krr_scan: + args: "--namespace devs-namespace" ## KRR args here + sinks: + - "main_slack_sink" # slack sink you want to send the report to here +``` + +3. Do a Helm upgrade to apply the new values: `helm upgrade robusta robusta/robusta --values=generated_values.yaml --set clusterName=<YOUR_CLUSTER_NAME>` + +<p align="right">(<a href="#readme-top">back to top</a>)</p> + <!-- GETTING STARTED --> ## Getting Started @@ -393,4 +427,5 @@ Project Link: [https://github.com/robusta-dev/krr](https://github.com/robusta-de [linkedin-shield]: https://img.shields.io/badge/-LinkedIn-black.svg?style=for-the-badge&logo=linkedin&colorB=555 [linkedin-url]: https://linkedin.com/in/othneildrew [product-screenshot]: images/screenshot.jpeg +[slack-screenshot]: images/krr_slack_example.png [ui-screenshot]: images/ui_screenshot.jpeg diff --git a/examples/custom_formatter.py b/examples/custom_formatter.py index d3cb98d..f246404 100644 --- a/examples/custom_formatter.py +++ b/examples/custom_formatter.py @@ -3,21 +3,17 @@ from __future__ import annotations import robusta_krr -from robusta_krr.api.formatters import BaseFormatter +from robusta_krr.api import formatters from robusta_krr.api.models import Result -class CustomFormatter(BaseFormatter): - # This is the name that will be used to reference the formatter in the CLI - __display_name__ = "my_formatter" - - # This will pass the result to Rich Console for formatting. - # By default, the result is passed to `print` function. - # See https://rich.readthedocs.io/en/latest/ for more info - __rich_console__ = True - - def format(self, result: Result) -> str: - return "Custom formatter" +# This is a custom formatter +# It will be available to the CLI as `my_formatter` +# Rich console will be enabled in this case, so the output will be colored and formatted +@formatters.register(rich_console=True) +def my_formatter(result: Result) -> str: + # Return custom formatter + return "Custom formatter" # Running this file will register the formatter and make it available to the CLI diff --git a/images/krr_slack_example.png b/images/krr_slack_example.png Binary files differnew file mode 100644 index 0000000..f7ce498 --- /dev/null +++ b/images/krr_slack_example.png diff --git a/pyproject.toml b/pyproject.toml index fe9317c..d6db5aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "robusta-krr" -version = "1.0.0" +version = "1.1.1" description = "Robusta's Resource Recommendation engine for Kubernetes" authors = ["Павел Жуков <33721692+LeaveMyYard@users.noreply.github.com>"] license = "MIT" diff --git a/robusta_krr/__init__.py b/robusta_krr/__init__.py index 3417aa1..cc72604 100644 --- a/robusta_krr/__init__.py +++ b/robusta_krr/__init__.py @@ -1,4 +1,4 @@ from .main import run -__version__ = "1.0.0" +__version__ = "1.1.1" __all__ = ["run", "__version__"] diff --git a/robusta_krr/api/formatters.py b/robusta_krr/api/formatters.py index 3bd6927..d942efe 100644 --- a/robusta_krr/api/formatters.py +++ b/robusta_krr/api/formatters.py @@ -1,3 +1,3 @@ -from robusta_krr.core.abstract.formatters import BaseFormatter +from robusta_krr.core.abstract.formatters import register, find, list_available -__all__ = ["BaseFormatter"] +__all__ = ["register", "find", "list_available"] diff --git a/robusta_krr/api/models.py b/robusta_krr/api/models.py index 5340f29..537168a 100644 --- a/robusta_krr/api/models.py +++ b/robusta_krr/api/models.py @@ -1,6 +1,6 @@ -from robusta_krr.core.abstract.strategies import HistoryData, ResourceRecommendation, RunResult +from robusta_krr.core.abstract.strategies import HistoryData, ResourceHistoryData, ResourceRecommendation, RunResult from robusta_krr.core.models.allocations import RecommendationValue, ResourceAllocations, ResourceType -from robusta_krr.core.models.objects import K8sObjectData +from robusta_krr.core.models.objects import K8sObjectData, PodData from robusta_krr.core.models.result import ResourceScan, Result from robusta_krr.core.models.severity import Severity, register_severity_calculator @@ -9,12 +9,13 @@ __all__ = [ "ResourceAllocations", "RecommendationValue", "K8sObjectData", + "PodData", "Result", "Severity", "register_severity_calculator", - "bind_calculator", "ResourceScan", "ResourceRecommendation", "HistoryData", + "ResourceHistoryData", "RunResult", ] diff --git a/robusta_krr/core/abstract/formatters.py b/robusta_krr/core/abstract/formatters.py index a38e9f3..b47d1e1 100644 --- a/robusta_krr/core/abstract/formatters.py +++ b/robusta_krr/core/abstract/formatters.py @@ -1,62 +1,80 @@ from __future__ import annotations -import abc -import os -from typing import TYPE_CHECKING, Any, TypeVar +from typing import Any, Optional, Callable -from robusta_krr.utils.display_name import add_display_name +from robusta_krr.core.models.result import Result -if TYPE_CHECKING: - from robusta_krr.core.models.result import Result +FormatterFunc = Callable[[Result], Any] -DEFAULT_FORMATTERS_PATH = os.path.join(os.path.dirname(__file__), "formatters") +FORMATTERS_REGISTRY: dict[str, FormatterFunc] = {} -Self = TypeVar("Self", bound="BaseFormatter") +# NOTE: Here asterisk is used to make the argument `rich_console` keyword-only +# This is done to avoid the following usage, where it is unclear what the boolean value is for: +# @register("My Formatter", True) +# def my_formatter(result: Result) -> str: +# return "My formatter" +# +# Instead, the following usage is enforced: +# @register("My Formatter", rich_console=True) +# def my_formatter(result: Result) -> str: +# return "My formatter" +def register(display_name: Optional[str] = None, *, rich_console: bool = False) -> Callable[[FormatterFunc], FormatterFunc]: + """ + A decorator to register a formatter function. -@add_display_name(postfix="Formatter") -class BaseFormatter(abc.ABC): - """Base class for result formatters.""" + Args: + display_name (str, optional): The name to use for the formatter in the registry. + rich_console (bool): Whether or not the formatter is for a rich console. Defaults to False. - __display_name__: str - __rich_console__: bool = False + Returns: + Callable[[FormatterFunc], FormatterFunc]: The decorator function. + """ - def __str__(self) -> str: - return self.__display_name__.title() + def decorator(func: FormatterFunc) -> FormatterFunc: + name = display_name or func.__name__ - @abc.abstractmethod - def format(self, result: Result) -> Any: - """Format the result. + FORMATTERS_REGISTRY[name] = func - Args: - result: The result to format. + func.__display_name__ = name # type: ignore + func.__rich_console__ = rich_console # type: ignore - Returns: - The formatted result. - """ + return func - @classmethod - def get_all(cls: type[Self]) -> dict[str, type[Self]]: - """Get all available formatters.""" + return decorator - # NOTE: Load default formatters - from robusta_krr import formatters as _ # noqa: F401 - return {sub_cls.__display_name__.lower(): sub_cls for sub_cls in cls.__subclasses__()} +def find(name: str) -> FormatterFunc: + """ + Find a formatter by name in the registry. - @staticmethod - def find(name: str) -> type[BaseFormatter]: - """Get a strategy from its name.""" + Args: + name (str): The name of the formatter. - formatters = BaseFormatter.get_all() + Returns: + FormatterFunc: The formatter function. - l_name = name.lower() - if l_name in formatters: - return formatters[l_name] + Raises: + ValueError: If a formatter with the given name does not exist. + """ - raise ValueError(f"Unknown formatter name: {name}. Available formatters: {', '.join(formatters)}") + try: + return FORMATTERS_REGISTRY[name] + except KeyError as e: + raise ValueError(f"Formatter '{name}' not found") from e -__all__ = ["BaseFormatter"] +def list_available() -> list[str]: + """ + List available formatters in the registry. + + Returns: + list[str]: A list of the names of the available formatters. + """ + + return list(FORMATTERS_REGISTRY) + + +__all__ = ["register", "find"] diff --git a/robusta_krr/core/abstract/strategies.py b/robusta_krr/core/abstract/strategies.py index 4f70e89..b0d9bf4 100644 --- a/robusta_krr/core/abstract/strategies.py +++ b/robusta_krr/core/abstract/strategies.py @@ -10,12 +10,18 @@ import pydantic as pd from numpy.typing import NDArray from robusta_krr.core.models.result import K8sObjectData, Metric, ResourceType -from robusta_krr.utils.display_name import add_display_name +from robusta_krr.utils.display_name import display_name_property SelfRR = TypeVar("SelfRR", bound="ResourceRecommendation") class ResourceRecommendation(pd.BaseModel): + """A class to represent resource recommendation with optional request and limit values. + + The NaN values are used to represent undefined values: the strategy did not provide a recommendation for the resource. + None values are used to represent the strategy says that value should not be set. + """ + request: Optional[float] limit: Optional[float] @@ -25,6 +31,15 @@ class ResourceRecommendation(pd.BaseModel): class StrategySettings(pd.BaseModel): + """A class to represent strategy settings with configurable history and timeframe duration. + + It is used in CLI to generate the help, parameters and validate values. + Description is used to generate the help. + Other pydantic features can be used to validate the values. + + Nested classes are not supported here. + """ + history_duration: float = pd.Field( 24 * 7 * 2, ge=1, description="The duration of the history data to use (in hours)." ) @@ -39,14 +54,19 @@ class StrategySettings(pd.BaseModel): return datetime.timedelta(minutes=self.timeframe_duration) -_StrategySettings = TypeVar("_StrategySettings", bound=StrategySettings) - +# A type alias for a numpy array of shape (N, 2). ArrayNx2 = Annotated[NDArray[np.float64], Literal["N", 2]] class ResourceHistoryData(pd.BaseModel): + """A class to represent resource history data. + + metric is the metric information used to gather the history data. + data is a mapping from pod to a numpy array of time and value. + """ + metric: Metric - data: dict[str, ArrayNx2] # Mapping: pod -> (time, value) + data: dict[str, ArrayNx2] # Mapping: pod -> [(time, value)] class Config: arbitrary_types_allowed = True @@ -56,10 +76,30 @@ HistoryData = dict[ResourceType, ResourceHistoryData] RunResult = dict[ResourceType, ResourceRecommendation] SelfBS = TypeVar("SelfBS", bound="BaseStrategy") +_StrategySettings = TypeVar("_StrategySettings", bound=StrategySettings) -@add_display_name(postfix="Strategy") +# An abstract base class for strategy implementation. +# This class requires implementation of a 'run' method for calculating recommendation. +# Make a subclass if you want to create a concrete strategy. +@display_name_property(suffix="Strategy") class BaseStrategy(abc.ABC, Generic[_StrategySettings]): + """An abstract base class for strategy implementation. + + This class is generic, and requires a type for the settings. + This settings type will be used for the settings property of the strategy. + It will be used to generate CLI parameters for this strategy, validated by pydantic. + + This class requires implementation of a 'run' method for calculating recommendation. + Additionally, it provides a 'description' property for generating a description for the strategy. + Description property uses the docstring of the strategy class and the settings of the strategy. + + The name of the strategy is the name of the class in lowercase, without the 'Strategy' suffix, if exists. + If you want to change the name of the strategy, you can change the __display_name__ attribute. + + The strategy will automatically be registered in the strategy registry using __subclasses__ mechanism. + """ + __display_name__: str settings: _StrategySettings @@ -83,27 +123,29 @@ class BaseStrategy(abc.ABC, Generic[_StrategySettings]): return f"[b]{self} Strategy[/b]\n\n" + dedent(self.__doc__.format_map(self.settings.dict())).strip() + # Abstract method that needs to be implemented by subclass. + # This method is intended to calculate resource recommendation based on history data and kubernetes object data. @abc.abstractmethod def run(self, history_data: HistoryData, object_data: K8sObjectData) -> RunResult: - """Run the strategy to calculate the recommendation""" + pass + # This method is intended to return a strategy by its name. @classmethod def find(cls: type[SelfBS], name: str) -> type[SelfBS]: - """Get a strategy from its name.""" - strategies = cls.get_all() if name.lower() in strategies: return strategies[name.lower()] raise ValueError(f"Unknown strategy name: {name}. Available strategies: {', '.join(strategies)}") + # This method is intended to return all the available strategies. @classmethod def get_all(cls: type[SelfBS]) -> dict[str, type[SelfBS]]: - # NOTE: Load default formatters from robusta_krr import strategies as _ # noqa: F401 return {sub_cls.__display_name__.lower(): sub_cls for sub_cls in cls.__subclasses__()} + # This method is intended to return the type of settings used in strategy. @classmethod def get_settings_type(cls) -> type[StrategySettings]: return get_args(cls.__orig_bases__[0])[0] # type: ignore @@ -111,6 +153,7 @@ class BaseStrategy(abc.ABC, Generic[_StrategySettings]): AnyStrategy = BaseStrategy[StrategySettings] + __all__ = [ "AnyStrategy", "BaseStrategy", diff --git a/robusta_krr/core/integrations/prometheus/loader.py b/robusta_krr/core/integrations/prometheus/loader.py index 1f0c14e..344e819 100644 --- a/robusta_krr/core/integrations/prometheus/loader.py +++ b/robusta_krr/core/integrations/prometheus/loader.py @@ -20,7 +20,20 @@ from .metrics import BaseMetricLoader class PrometheusDiscovery(ServiceDiscovery): + """ + Service discovery for Prometheus. + """ + def find_prometheus_url(self, *, api_client: Optional[ApiClient] = None) -> Optional[str]: + """ + Finds the Prometheus URL using selectors. + + Args: + api_client (Optional[ApiClient]): A Kubernetes API client. Defaults to None. + + Returns: + Optional[str]: The discovered Prometheus URL, or None if not found. + """ return super().find_url( selectors=[ "app=kube-prometheus-stack-prometheus", @@ -30,16 +43,25 @@ class PrometheusDiscovery(ServiceDiscovery): "app=prometheus-msteams", "app=rancher-monitoring-prometheus", "app=prometheus-prometheus", + "app.kubernetes.io/name=vmsingle", ], api_client=api_client, ) class PrometheusNotFound(Exception): + """ + An exception raised when Prometheus is not found. + """ + pass class CustomPrometheusConnect(PrometheusConnect): + """ + Custom PrometheusConnect class to handle retries. + """ + @no_type_check def __init__( self, @@ -55,12 +77,24 @@ class CustomPrometheusConnect(PrometheusConnect): class PrometheusLoader(Configurable): + """ + A loader class for fetching metrics from Prometheus. + """ + def __init__( self, config: Config, *, cluster: Optional[str] = None, ) -> None: + """ + Initializes the Prometheus Loader. + + Args: + config (Config): The configuration object. + cluster (Optional[str]): The name of the cluster. Defaults to None. + """ + super().__init__(config=config) self.info(f"Connecting to Prometheus for {cluster or 'default'} cluster") @@ -93,6 +127,13 @@ class PrometheusLoader(Configurable): self.info(f"Prometheus connected successfully for {cluster or 'default'} cluster") def _check_prometheus_connection(self): + """ + Checks the connection to Prometheus. + + Raises: + PrometheusNotFound: If the connection to Prometheus cannot be established. + """ + try: response = self.prometheus._session.get( f"{self.prometheus.url}/api/v1/query", @@ -115,6 +156,19 @@ class PrometheusLoader(Configurable): *, step: datetime.timedelta = datetime.timedelta(minutes=30), ) -> ResourceHistoryData: + """ + Gathers data from Prometheus for a specified object and resource. + + Args: + object (K8sObjectData): The Kubernetes object. + resource (ResourceType): The resource type. + period (datetime.timedelta): The time period for which to gather data. + step (datetime.timedelta, optional): The time step between data points. Defaults to 30 minutes. + + Returns: + ResourceHistoryData: The gathered resource history data. + """ + self.debug(f"Gathering data for {object} and {resource}") await self.add_historic_pods(object, period) @@ -124,12 +178,20 @@ class PrometheusLoader(Configurable): return await metric_loader.load_data(object, period, step) async def add_historic_pods(self, object: K8sObjectData, period: datetime.timedelta) -> None: - """Find pods that were already deleted, but still have some metrics in Prometheus""" + """ + Finds pods that have been deleted but still have some metrics in Prometheus. + + Args: + object (K8sObjectData): The Kubernetes object. + period (datetime.timedelta): The time period for which to gather data. + """ if len(object.pods) == 0: return - - period_literal = f"{int(period.total_seconds()) // 60 // 24}d" + + # Prometheus limit, the max can be 32 days + days_literal = min(int(period.total_seconds()) // 60 // 24, 32) + period_literal = f"{days_literal}d" owner = await asyncio.to_thread( self.prometheus.custom_query, query=f'kube_pod_owner{{pod="{next(iter(object.pods)).name}"}}[{period_literal}]', diff --git a/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py b/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py index 9ede988..6a4d964 100644 --- a/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py +++ b/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py @@ -22,7 +22,6 @@ class BaseFilteredMetricLoader(BaseMetricLoader): return series["metric"][label] return None - # TODO: Rework this, as now our query can return multiple metrics for different pods @staticmethod def filter_prom_jobs_results( series_list_result: list[PrometheusSeries], diff --git a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py index 2751f32..a0731c5 100644 --- a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py +++ b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py @@ -15,22 +15,59 @@ from robusta_krr.utils.configurable import Configurable if TYPE_CHECKING: from ..loader import CustomPrometheusConnect +# A registry of metrics that can be used to fetch a corresponding metric loader. REGISTERED_METRICS: dict[str, type[BaseMetricLoader]] = {} class BaseMetricLoader(Configurable, abc.ABC): + """ + Base class for all metric loaders. + + Metric loaders are used to load metrics from a specified source (like Prometheus in this case). + """ + def __init__(self, config: Config, prometheus: CustomPrometheusConnect) -> None: super().__init__(config) self.prometheus = prometheus @abc.abstractmethod def get_query(self, object: K8sObjectData) -> str: - ... + """ + This method should be implemented by all subclasses to provide a query string to fetch metrics. + + Args: + object (K8sObjectData): The object for which metrics need to be fetched. + + Returns: + str: The query string. + """ + + pass def _step_to_string(self, step: datetime.timedelta) -> str: + """ + Converts step in datetime.timedelta format to a string format used by Prometheus. + + Args: + step (datetime.timedelta): Step size in datetime.timedelta format. + + Returns: + str: Step size in string format used by Prometheus. + """ + return f"{int(step.total_seconds()) // 60}m" async def query_prometheus(self, metric: Metric) -> list[dict]: + """ + Asynchronous method that queries Prometheus to fetch metrics. + + Args: + metric (Metric): An instance of the Metric class specifying what metrics to fetch. + + Returns: + list[dict]: A list of dictionary where each dictionary represents metrics for a pod. + """ + return await asyncio.to_thread( self.prometheus.custom_query_range, query=metric.query, @@ -42,8 +79,20 @@ class BaseMetricLoader(Configurable, abc.ABC): async def load_data( self, object: K8sObjectData, period: datetime.timedelta, step: datetime.timedelta ) -> ResourceHistoryData: + """ + Asynchronous method that loads metric data for a specific object. + + Args: + object (K8sObjectData): The object for which metrics need to be loaded. + period (datetime.timedelta): The time period for which metrics need to be loaded. + step (datetime.timedelta): The time interval between successive metric values. + + Returns: + ResourceHistoryData: An instance of the ResourceHistoryData class representing the loaded metrics. + """ + query = self.get_query(object) - end_time = datetime.datetime.now() + end_time = datetime.datetime.now().astimezone() metric = Metric( query=query, start_time=end_time - period, @@ -65,6 +114,19 @@ class BaseMetricLoader(Configurable, abc.ABC): @staticmethod def get_by_resource(resource: str) -> type[BaseMetricLoader]: + """ + Fetches the metric loader corresponding to the specified resource. + + Args: + resource (str): The name of the resource. + + Returns: + type[BaseMetricLoader]: The class of the metric loader corresponding to the resource. + + Raises: + KeyError: If the specified resource is not registered. + """ + try: return REGISTERED_METRICS[resource] except KeyError as e: @@ -75,6 +137,16 @@ Self = TypeVar("Self", bound=BaseMetricLoader) def bind_metric(resource: str) -> Callable[[type[Self]], type[Self]]: + """ + A decorator that binds a metric loader to a resource. + + Args: + resource (str): The name of the resource. + + Returns: + Callable[[type[Self]], type[Self]]: The decorator that does the binding. + """ + def decorator(cls: type[Self]) -> type[Self]: REGISTERED_METRICS[resource] = cls return cls diff --git a/robusta_krr/core/models/config.py b/robusta_krr/core/models/config.py index ff147af..6a79a19 100644 --- a/robusta_krr/core/models/config.py +++ b/robusta_krr/core/models/config.py @@ -4,7 +4,7 @@ import pydantic as pd from kubernetes import config from kubernetes.config.config_exception import ConfigException -from robusta_krr.core.abstract.formatters import BaseFormatter +from robusta_krr.core.abstract import formatters from robusta_krr.core.abstract.strategies import AnyStrategy, BaseStrategy try: @@ -44,8 +44,8 @@ class Config(pd.BaseSettings): other_args: dict[str, Any] @property - def Formatter(self) -> type[BaseFormatter]: - return BaseFormatter.find(self.format) + def Formatter(self) -> formatters.FormatterFunc: + return formatters.find(self.format) @pd.validator("namespaces") def validate_namespaces(cls, v: Union[list[str], Literal["*"]]) -> Union[list[str], Literal["*"]]: @@ -66,7 +66,7 @@ class Config(pd.BaseSettings): @pd.validator("format") def validate_format(cls, v: str) -> str: - BaseFormatter.find(v) # NOTE: raises if strategy is not found + formatters.find(v) # NOTE: raises if strategy is not found return v @property diff --git a/robusta_krr/core/models/objects.py b/robusta_krr/core/models/objects.py index 21f5d61..84e468e 100644 --- a/robusta_krr/core/models/objects.py +++ b/robusta_krr/core/models/objects.py @@ -1,5 +1,3 @@ -from typing import Optional - import pydantic as pd from robusta_krr.core.models.allocations import ResourceAllocations @@ -14,12 +12,12 @@ class PodData(pd.BaseModel): class K8sObjectData(pd.BaseModel): - cluster: Optional[str] + cluster: str name: str container: str pods: list[PodData] namespace: str - kind: Optional[str] + kind: str allocations: ResourceAllocations def __str__(self) -> str: diff --git a/robusta_krr/core/models/result.py b/robusta_krr/core/models/result.py index 46512bb..a72ff21 100644 --- a/robusta_krr/core/models/result.py +++ b/robusta_krr/core/models/result.py @@ -5,7 +5,7 @@ from typing import Any, Optional, Union import pydantic as pd -from robusta_krr.core.abstract.formatters import BaseFormatter +from robusta_krr.core.abstract import formatters from robusta_krr.core.models.allocations import RecommendationValue, ResourceAllocations, ResourceType from robusta_krr.core.models.objects import K8sObjectData from robusta_krr.core.models.severity import Severity @@ -75,7 +75,7 @@ class Result(pd.BaseModel): super().__init__(*args, **kwargs) self.score = self.__calculate_score() - def format(self, formatter: Union[type[BaseFormatter], str], **kwargs: Any) -> Any: + def format(self, formatter: Union[formatters.FormatterFunc, str]) -> Any: """Format the result. Args: @@ -85,9 +85,8 @@ class Result(pd.BaseModel): The formatted result. """ - FormatterType = BaseFormatter.find(formatter) if isinstance(formatter, str) else formatter - _formatter = FormatterType(**kwargs) - return _formatter.format(self) + formatter = formatters.find(formatter) if isinstance(formatter, str) else formatter + return formatter(self) @staticmethod def __scan_cost(scan: ResourceScan) -> float: diff --git a/robusta_krr/formatters/__init__.py b/robusta_krr/formatters/__init__.py index 0fc1c80..325cf01 100644 --- a/robusta_krr/formatters/__init__.py +++ b/robusta_krr/formatters/__init__.py @@ -1,4 +1,4 @@ -from .json import JSONFormatter -from .pprint import PPrintFormatter -from .table import TableFormatter -from .yaml import YAMLFormatter +from .json import json +from .pprint import pprint +from .table import table +from .yaml import yaml diff --git a/robusta_krr/formatters/json.py b/robusta_krr/formatters/json.py index 2c3a51e..c391da2 100644 --- a/robusta_krr/formatters/json.py +++ b/robusta_krr/formatters/json.py @@ -1,21 +1,7 @@ -from __future__ import annotations - -from robusta_krr.core.abstract.formatters import BaseFormatter +from robusta_krr.core.abstract import formatters from robusta_krr.core.models.result import Result -class JSONFormatter(BaseFormatter): - """Formatter for JSON output.""" - - __display_name__ = "json" - - def format(self, result: Result) -> str: - """Format the result as JSON. - - :param result: The results to format. - :type result: :class:`core.result.Result` - :returns: The formatted results. - :rtype: str - """ - - return result.json(indent=2) +@formatters.register() +def json(result: Result) -> str: + return result.json(indent=2) diff --git a/robusta_krr/formatters/pprint.py b/robusta_krr/formatters/pprint.py index bdfcc4c..9be637c 100644 --- a/robusta_krr/formatters/pprint.py +++ b/robusta_krr/formatters/pprint.py @@ -1,23 +1,9 @@ -from __future__ import annotations - from pprint import pformat -from robusta_krr.core.abstract.formatters import BaseFormatter +from robusta_krr.core.abstract import formatters from robusta_krr.core.models.result import Result -class PPrintFormatter(BaseFormatter): - """Formatter for object output with python's pprint module.""" - - __display_name__ = "pprint" - - def format(self, result: Result) -> str: - """Format the result using pprint.pformat(...) - - :param result: The results to format. - :type result: :class:`core.result.Result` - :returns: The formatted results. - :rtype: str - """ - - return pformat(result.dict()) +@formatters.register() +def pprint(result: Result) -> str: + return pformat(result.dict()) diff --git a/robusta_krr/formatters/table.py b/robusta_krr/formatters/table.py index 6faef56..c510fc2 100644 --- a/robusta_krr/formatters/table.py +++ b/robusta_krr/formatters/table.py @@ -1,10 +1,8 @@ -from __future__ import annotations - import itertools from rich.table import Table -from robusta_krr.core.abstract.formatters import BaseFormatter +from robusta_krr.core.abstract import formatters from robusta_krr.core.models.allocations import RecommendationValue from robusta_krr.core.models.result import ResourceScan, ResourceType, Result from robusta_krr.utils import resource_units @@ -13,87 +11,78 @@ NONE_LITERAL = "unset" NAN_LITERAL = "?" -class TableFormatter(BaseFormatter): - """Formatter for text output.""" - - __display_name__ = "table" - __rich_console__ = True - - def _format(self, value: RecommendationValue) -> str: - if value is None: - return NONE_LITERAL - elif isinstance(value, str): - return NAN_LITERAL - else: - return resource_units.format(value) - - def _format_request_str(self, item: ResourceScan, resource: ResourceType, selector: str) -> str: - allocated = getattr(item.object.allocations, selector)[resource] - recommended = getattr(item.recommended, selector)[resource] - severity = recommended.severity - - return ( - f"[{severity.color}]" - + self._format(allocated) - + " -> " - + self._format(recommended.value) - + f"[/{severity.color}]" - ) - - def format(self, result: Result) -> Table: - """Format the result as text. - - :param result: The result to format. - :type result: :class:`core.result.Result` - :returns: The formatted results. - :rtype: str - """ - - table = Table( - show_header=True, - header_style="bold magenta", - title=f"\n{result.description}\n" if result.description else None, - title_justify="left", - title_style="", - caption=f"{result.score} points - {result.score_letter}", - ) - - table.add_column("Number", justify="right", no_wrap=True) - table.add_column("Cluster", style="cyan") - table.add_column("Namespace", style="cyan") - table.add_column("Name", style="cyan") - table.add_column("Pods", style="cyan") - table.add_column("Old Pods", style="cyan") - table.add_column("Type", style="cyan") - table.add_column("Container", style="cyan") - for resource in ResourceType: - table.add_column(f"{resource.name} Requests") - table.add_column(f"{resource.name} Limits") - - for _, group in itertools.groupby( - enumerate(result.scans), key=lambda x: (x[1].object.cluster, x[1].object.namespace, x[1].object.name) - ): - group_items = list(group) - - for j, (i, item) in enumerate(group_items): - last_row = j == len(group_items) - 1 - full_info_row = j == 0 - - table.add_row( - f"[{item.severity.color}]{i + 1}.[/{item.severity.color}]", - item.object.cluster if full_info_row else "", - item.object.namespace if full_info_row else "", - item.object.name if full_info_row else "", - f"{item.object.current_pods_count}" if full_info_row else "", - f"{item.object.deleted_pods_count}" if full_info_row else "", - item.object.kind if full_info_row else "", - item.object.container, - *[ - self._format_request_str(item, resource, selector) - for resource in ResourceType - for selector in ["requests", "limits"] - ], - end_section=last_row, - ) - - return table +def _format(value: RecommendationValue) -> str: + if value is None: + return NONE_LITERAL + elif isinstance(value, str): + return NAN_LITERAL + else: + return resource_units.format(value) + + +def _format_request_str(item: ResourceScan, resource: ResourceType, selector: str) -> str: + allocated = getattr(item.object.allocations, selector)[resource] + recommended = getattr(item.recommended, selector)[resource] + severity = recommended.severity + + return f"[{severity.color}]" + _format(allocated) + " -> " + _format(recommended.value) + f"[/{severity.color}]" + + +@formatters.register(rich_console=True) +def table(result: Result) -> Table: + """Format the result as text. + + :param result: The result to format. + :type result: :class:`core.result.Result` + :returns: The formatted results. + :rtype: str + """ + + table = Table( + show_header=True, + header_style="bold magenta", + title=f"\n{result.description}\n" if result.description else None, + title_justify="left", + title_style="", + caption=f"{result.score} points - {result.score_letter}", + ) + + table.add_column("Number", justify="right", no_wrap=True) + table.add_column("Cluster", style="cyan") + table.add_column("Namespace", style="cyan") + table.add_column("Name", style="cyan") + table.add_column("Pods", style="cyan") + table.add_column("Old Pods", style="cyan") + table.add_column("Type", style="cyan") + table.add_column("Container", style="cyan") + for resource in ResourceType: + table.add_column(f"{resource.name} Requests") + table.add_column(f"{resource.name} Limits") + + for _, group in itertools.groupby( + enumerate(result.scans), key=lambda x: (x[1].object.cluster, x[1].object.namespace, x[1].object.name) + ): + group_items = list(group) + + for j, (i, item) in enumerate(group_items): + last_row = j == len(group_items) - 1 + full_info_row = j == 0 + + table.add_row( + f"[{item.severity.color}]{i + 1}.[/{item.severity.color}]", + item.object.cluster if full_info_row else "", + item.object.namespace if full_info_row else "", + item.object.name if full_info_row else "", + f"{item.object.current_pods_count}" if full_info_row else "", + f"{item.object.deleted_pods_count}" if full_info_row else "", + item.object.kind if full_info_row else "", + item.object.container, + *[ + _format_request_str(item, resource, selector) + for resource in ResourceType + for selector in ["requests", "limits"] + ], + end_section=last_row, + ) + + return table diff --git a/robusta_krr/formatters/yaml.py b/robusta_krr/formatters/yaml.py index c494bc1..37a030c 100644 --- a/robusta_krr/formatters/yaml.py +++ b/robusta_krr/formatters/yaml.py @@ -1,24 +1,11 @@ -from __future__ import annotations - import json -import yaml +import yaml as yaml_module -from robusta_krr.core.abstract.formatters import BaseFormatter +from robusta_krr.core.abstract import formatters from robusta_krr.core.models.result import Result -class YAMLFormatter(BaseFormatter): - """Formatter for YAML output.""" - - __display_name__ = "yaml" - - def format(self, result: Result) -> str: - """Format the result as YAML. - - :param result: The results to format. - :type result: :class:`core.result.Result` - :returns: The formatted results. - :rtype: str - """ - return yaml.dump(json.loads(result.json()), sort_keys=False) +@formatters.register() +def yaml(result: Result) -> str: + return yaml_module.dump(json.loads(result.json()), sort_keys=False) diff --git a/robusta_krr/main.py b/robusta_krr/main.py index fdd654d..1734788 100644 --- a/robusta_krr/main.py +++ b/robusta_krr/main.py @@ -9,8 +9,9 @@ from uuid import UUID import typer import urllib3 -from robusta_krr.core.abstract.formatters import BaseFormatter +from robusta_krr import formatters as concrete_formatters # noqa: F401 from robusta_krr.core.abstract.strategies import AnyStrategy, BaseStrategy +from robusta_krr.core.abstract import formatters from robusta_krr.core.models.config import Config from robusta_krr.core.runner import Runner from robusta_krr.utils.version import get_version @@ -115,7 +116,7 @@ def load_commands() -> None: f"'{field_name}': {field_name}" for field_name in strategy_type.get_settings_type().__fields__ ) + "}", - formatters=", ".join(BaseFormatter.get_all()), + formatters=", ".join(formatters.list_available()), ), globals() | {strategy.__name__: strategy for strategy in AnyStrategy.get_all().values()} # Defined strategies diff --git a/robusta_krr/utils/configurable.py b/robusta_krr/utils/configurable.py index 4777a9c..c2b2465 100644 --- a/robusta_krr/utils/configurable.py +++ b/robusta_krr/utils/configurable.py @@ -10,7 +10,9 @@ from robusta_krr.core.models.config import Config class Configurable(abc.ABC): """ A class that can be configured with a Config object. - Opens the possibility to use echo and debug methods + Opens the possibility to use custom logging methods, that can be configured with the Config object. + + Also makes a `console` attribute available, which is a rich console. """ def __init__(self, config: Config) -> None: diff --git a/robusta_krr/utils/display_name.py b/robusta_krr/utils/display_name.py index 9082fc0..69541ae 100644 --- a/robusta_krr/utils/display_name.py +++ b/robusta_krr/utils/display_name.py @@ -1,16 +1,27 @@ -from typing import Callable, TypeVar +from typing import Callable, TypeVar, Any _T = TypeVar("_T") -def add_display_name(*, postfix: str) -> Callable[[type[_T]], type[_T]]: - """Add a decorator factory to add __display_name__ property to the class.""" +def display_name_property(*, suffix: str) -> Callable[[type[_T]], type[_T]]: + """Add a decorator factory to add __display_name__ property to the class. + + It is a utility function for BaseStrategy. + It makes a __display_name__ property for the class, that uses the name of the class. + By default, it will remove the suffix from the name of the class. + For example, if the name of the class is 'MyStrategy', the __display_name__ property will be 'My'. + If the name of the class is 'Foo', the __display_name__ property will be 'Foo', because it does not end with 'Strategy'. + + If you then override the __display_name__ property, it will be used instead of the default one. + """ def decorator(cls: type[_T]) -> type[_T]: class DisplayNameProperty: - def __get__(self, instance, owner): - if owner.__name__.lower().endswith(postfix.lower()): - return owner.__name__[: -len(postfix)] + # This is a descriptor that returns the name of the class. + # It is used to generate the __display_name__ property. + def __get__(self, instance: Any, owner: type[_T]) -> str: + if owner.__name__.lower().endswith(suffix.lower()): + return owner.__name__[: -len(suffix)] return owner.__name__ diff --git a/robusta_krr/utils/progress_bar.py b/robusta_krr/utils/progress_bar.py index 808a380..32bf211 100644 --- a/robusta_krr/utils/progress_bar.py +++ b/robusta_krr/utils/progress_bar.py @@ -5,6 +5,13 @@ from robusta_krr.utils.configurable import Configurable class ProgressBar(Configurable): + """ + Progress bar for displaying progress of gathering recommendations. + + Use `ProgressBar` as a context manager to automatically handle the progress bar. + Use `progress` method to step the progress bar. + """ + def __init__(self, config: Config, **kwargs) -> None: super().__init__(config) self.show_bar = self.echo_active diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..0617ddc --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,81 @@ +import random +from datetime import datetime, timedelta +from unittest.mock import AsyncMock, PropertyMock, patch + +import numpy as np +import pytest + +from robusta_krr.api.models import K8sObjectData, PodData, ResourceAllocations, ResourceHistoryData + +TEST_OBJECT = K8sObjectData( + cluster="mock-cluster", + name="mock-object-1", + container="mock-container-1", + pods=[ + PodData(name="mock-pod-1", deleted=False), + PodData(name="mock-pod-2", deleted=False), + PodData(name="mock-pod-3", deleted=True), + ], + namespace="default", + kind="Deployment", + allocations=ResourceAllocations( + requests={"cpu": 1, "memory": 1}, # type: ignore + limits={"cpu": 2, "memory": 2}, # type: ignore + ), +) + + +@pytest.fixture(autouse=True, scope="session") +def mock_list_clusters(): + with patch( + "robusta_krr.core.integrations.kubernetes.KubernetesLoader.list_clusters", + new=AsyncMock(return_value=[TEST_OBJECT.cluster]), + ): + yield + + +@pytest.fixture(autouse=True, scope="session") +def mock_list_scannable_objects(): + with patch( + "robusta_krr.core.integrations.kubernetes.KubernetesLoader.list_scannable_objects", + new=AsyncMock(return_value=[TEST_OBJECT]), + ): + yield + + +@pytest.fixture(autouse=True, scope="session") +def mock_config_loaded(): + with patch("robusta_krr.core.models.config.Config.config_loaded", new_callable=PropertyMock) as mock_config: + mock_config.return_value = True + yield + + +@pytest.fixture(autouse=True, scope="session") +def mock_prometheus_loader(): + now = datetime.now() + start = now - timedelta(hours=1) + now_ts, start_ts = now.timestamp(), start.timestamp() + metric_points_data = np.array([(t, random.randrange(0, 100)) for t in np.linspace(start_ts, now_ts, 3600)]) + + with patch( + "robusta_krr.core.integrations.prometheus.loader.PrometheusLoader.gather_data", + new=AsyncMock( + return_value=ResourceHistoryData( + data={pod.name: metric_points_data for pod in TEST_OBJECT.pods}, + metric={ # type: ignore + "query": f"example_promql_metric{{pod_name=~\"{'|'.join(pod.name for pod in TEST_OBJECT.pods)}\"}}", + "start_time": start, + "end_time": now, + "step": "30s", + }, + ) + ), + ) as mock_prometheus_loader: + mock_prometheus_loader + yield + + +@pytest.fixture(autouse=True, scope="session") +def mock_prometheus_init(): + with patch("robusta_krr.core.integrations.prometheus.loader.PrometheusLoader.__init__", return_value=None): + yield diff --git a/tests/test_krr.py b/tests/test_krr.py index 23dab01..f013ee4 100644 --- a/tests/test_krr.py +++ b/tests/test_krr.py @@ -1,8 +1,3 @@ -""" - Test the krr command line interface and a general execution. - Requires a running kubernetes cluster with the kubectl command configured. -""" - import json import pytest @@ -36,14 +31,19 @@ def test_run(log_flag: str): @pytest.mark.parametrize("format", ["json", "yaml", "table", "pprint"]) def test_output_formats(format: str): - result = runner.invoke(app, [STRATEGY_NAME, "-q", "-f", format, "--namespace", "default"]) + result = runner.invoke(app, [STRATEGY_NAME, "-q", "-f", format]) try: assert result.exit_code == 0, result.exc_info except AssertionError as e: raise e from result.exception - if format == "json": - assert json.loads(result.stdout), result.stdout - - if format == "yaml": - assert yaml.safe_load(result.stdout), result.stdout + try: + if format == "json": + json_output = json.loads(result.stdout) + assert json_output, result.stdout + assert len(json_output["scans"]) > 0, result.stdout + + if format == "yaml": + assert yaml.safe_load(result.stdout), result.stdout + except Exception as e: + raise Exception(result.stdout) from e |
