Writing Plugins

Even though infra-tester provides several inbuilt assertions for tests, there will always be a case where you'll have to test something that might not be possible with just the inbuilt set of assertions.

Since infrastructure varies a lot this is bound to happen, and a good way to tackle this is through plugins.

infra-tester provides a plugin system where you can create your custom assertions in Python and then consume them through YAML configuration. This way, plugin owners can create assertions that can then be shared or published so that others who may not be proficient in Python could use them by installing the plugin.

Let's see how the plugin system works.

Plugin Design

infra-tester chose Python 3.8 and above as the plugin language for the following reasons:

• Python is a relatively simple language
• Has an extensive standard library that could assist in implementing plugins
• Entry Point Spec which makes it an ideal ecosystem to write plugins
• The PIP package manager makes it easy to package and install plugins through the official PIP registry, a custom registry, or through a version control system or even install the plugin from a local file system

The below sequence diagram shows how infra-tester interacts with the plugins:

sequenceDiagram
    infra-tester->>plugin-host: Fetch available plugins
    plugin-host-->>infra-tester: <List of available plugins>

    infra-tester->>plugin-host: Run plugin validation
    plugin-host->>plugin: Calls `validate_assertion` function of the plugin
    plugin-->>plugin-host: Returns validation error message if any
    plugin-host-->>infra-tester: Returns validation result

    infra-tester->>plugin-host: Run plugin assertion
    plugin-host->>plugin: Calls `run_assertion` function of the plugin
    plugin-->>plugin-host: Returns error message if any
    plugin-host-->>infra-tester: Returns assertion result

    infra-tester->>plugin-host: Run plugin cleanup
    plugin-host->>plugin: Calls `cleanup` function of the plugin
    plugin-->>plugin-host: Returns error message if any
    plugin-host-->>infra-tester: Returns cleanup result

In the above diagram:

• plugin-host is a PIP package, which if installed enables plugin support for infra-tester.
• plugin is a PIP package that can be used by infra-tester as a plugin.

infra-tester uses the plugin-host to communicate with the plugin.

As a plugin implementor, you only have to be concerned about the plugin package.

Plugin Host

You can install the plugin host by installing the infra-tester-plugins PIP package.

Info

As of now, you can install the plugin host through PIP using the following command:

pip install "git+https://github.com/schrodinger/infra-tester.git#subdirectory=python-plugins/"

We are currently looking into adding this package to the official PIP registry.

Once the package is installed, infra-tester will automatically enable plugin support.

Plugins

Requirements

Before starting, ensure that infra-tester-plugins package is installed. This provides the BaseAssertionPlugin class, which needs to be inherited to implement the plugin.

As a best practice, it's recommended to follow a standard directory structure when creating the Python package. Generally, it's recommended to follow the below directory structure:

.
└── <package-directory-name>/
    ├── src/
    │   ├── <module1>/
    │   │   ├── __init__.py
    │   │   └── ...
    │   └── <module2>/
    │       ├── __init__.py
    │       └── ...
    ├── pyproject.toml
    └── setup.py

The BaseAssertionPlugin Class

infra-tester provides the BaseAssertionPlugin base class which can be used to implement the core of the plugin logic.

from infra_tester_plugins import BaseAssertionPlugin

class CustomAssertionPlugin(BaseAssertionPlugin):
    def validate_inputs(self, inputs: dict): ...

    def run_assertion(self, inputs: dict, state: dict): ...

    def cleanup(self, inputs: dict, state: dict): ... # optional


def load_plugin() # A function that'll be referenced in pyproject.toml.
    # This function must return an instance of the assertion implementation.

    return CustomAssertionPlugin()

def validate_inputs(self, inputs: dict) -> Union[str, None]:

Validate the inputs provided to the plugin. The inputs will be in the form of a dictionary with the key being the name of the input and the value being the value of the input.

This method should return None if the inputs are valid. Otherwise, it should return a string describing the error.

Any exception thrown by this method will be treated as an implementation error and will be logged as such.

ArgumentsExceptionsReturn Value

Name	Description	Type
inputs	Python dictionary containing the inputs provided to the assertion	Dictionary

Exception	Description
NotImplementedError	If the plugin does not implement this method.

Type	Description
Union[str, None]	If the inputs are valid, return None. Otherwise, return a string describing the error.

def run_assertion(self, inputs: dict, state: dict) -> Union[str, None]:

This method should contain the logic to run the assertion and return the result.

If the assertion fails, this method should return a string describing the error. Otherwise, it should return None.

This method should not raise any exception. Any exception thrown by this method will be treated as an implementation error and will be logged as such.

ArgumentsExceptionsReturn Value

Name	Description	Type
inputs	Python dictionary containing the inputs provided to the assertion	Dictionary
state	Python dictionary containing the Terraform state information	Dictionary

Exception	Description
NotImplementedError	If the plugin does not implement this method.

Type	Description
Union[str, None]	None if the assertion passes. Otherwise return a string describing why the assertion failed.

def cleanup(self, inputs: dict, state: dict) -> None:

Cleanup any resources used by the plugin. This method will be called after the plugin has been run regardless of whether the run was successful or not. It is optional and does not need to be implemented if there's nothing to clean up.

If this method is implemented, it should be idempotent, i.e, cleanup should always have the same end result regardless of how many times it is called. Cleanup is considered successful if it does not raise any exception.

Arguments

Name	Description	Type
inputs	Python dictionary containing the inputs provided to the assertion	Dictionary
state	Python dictionary containing the Terraform state information	Dictionary

pyproject.toml and setup.py

The plugin packages need to provide pyproject.toml and setup.py files for the following reasons:

• Make the plugin installable through PIP package manager.
• Register the package as an infra-tester plugin.

See project.toml documentation to create a pyproject.toml file for your plugin.

Apart from this, one important part of creating pyproject.toml file is to register the package as a plugin by adding it to the infra-tester.assertion entry points group. This can be done by defining a [project.entry-points."infra_tester.assertion"] section in the pyproject.toml file, like so:

pyproject.tomlExample

[project.entry-points."infra_tester.assertion"]
<UniqueAssertionName1> = "<package_or_module1>:load_plugin"
<UniqueAssertionName1> = "<package_or_module1>:load_plugin"
...

[project.entry-points."infra_tester.assertion"]
ExampleAssertion = "example:load_plugin"
URLReachable = "url_reachable:load_plugin"

Info

A single python package can register multiple plugins (assertions). This makes it easy to create for example an "infra-tester-network-assertion" which provides several different network assertions as plugins.

It's strongly recommended to create a setup.py for the project as well to support older PIP versions too. You can use a shim in setup.py to not have duplicated configuration like so:

import setuptools

if __name__ == "__main__":
    setuptools.setup()

Example Plugin Package

This section shows an example package that provides two plugins:

• "ExampleAssertion" that just prints the inputs and state in the test logs.
• "URLReachable" that checks whether a URL is reachable.

ExampleAssertionURLReachablepyproject.tomlUsage in .infra-tester-config.yamlPackage Directory Structure

from typing import Dict

from infra_tester_plugins import BaseAssertionPlugin


class ExampleAssertionPlugin(BaseAssertionPlugin):
    def validate_inputs(self, inputs: Dict[str, object]):
        print("Running validate_inputs from ExampleAssertionPlugin")
        print("Inputs:", inputs)

        if (
            "should_fail_validation" in inputs
            and inputs["should_fail_validation"]
        ):
            return "This is a validation error message."

        return None

    def run_assertion(
        self, inputs: Dict[str, object], state: Dict[str, object]
    ):
        print("Running run_assertion from ExampleAssertionPlugin")
        print("Inputs:", inputs)
        print("State:", state)

        if (
            "should_error" in inputs
            and inputs["should_error"]
        ):
            message = "This is an assertion error message."
            if "custom_message" in state and state["custom_message"]:
                message += f" Plus here's a custom message: \
                            {state['custom_message']}"

            return message

        return None

    def cleanup(self, inputs: Dict[str, object], state: Dict[str, object]):
        print("Running cleanup from ExampleAssertionPlugin")
        print("Inputs:", inputs)
        print("State:", state)


def load_plugin():
    return ExampleAssertionPlugin()

from typing import Dict
from urllib import error, request
from urllib.parse import urlparse

from infra_tester_plugins import BaseAssertionPlugin


class URLReachableAssertionPlugin(BaseAssertionPlugin):
    FIELD_URL = "url"
    FIELD_STATUS_CODE = "status_code"
    FIELD_FROM_OUTPUTS = "from_outputs"

    @classmethod
    def validate_url(cls, plugin_inputs: Dict[str, object]):
        if cls.FIELD_URL not in plugin_inputs:
            return "url is a required input."

        if cls.FIELD_FROM_OUTPUTS in plugin_inputs:
            from_outputs = plugin_inputs[cls.FIELD_FROM_OUTPUTS]

            if not isinstance(from_outputs, bool):
                return "from_outputs must be a boolean."

            if from_outputs:
                return None

        url = plugin_inputs[cls.FIELD_URL]

        if not isinstance(url, str):
            return "url must be a string."

        try:
            parse_res = urlparse(url)

            if not parse_res.scheme:
                return "url must have a scheme."

            if not parse_res.netloc:
                return "url must have a valid network location."
        except ValueError:
            return "url must be a valid."

        return None

    @classmethod
    def validate_status_code(cls, plugin_inputs: Dict[str, object]):
        # status_code is optional
        if cls.FIELD_STATUS_CODE not in plugin_inputs:
            return None

        status_code = plugin_inputs[cls.FIELD_STATUS_CODE]

        if not isinstance(status_code, int):
            return "status_code must be an integer."

        # Not being strict about the status code range here.
        if status_code < 0:
            return (
                f"status_code must be valid. "
                f"Received {status_code} of type {type(status_code)}."
            )

        return None

    def validate_inputs(self, inputs: Dict[str, object]):
        print("Running validate_inputs for URLReachable")

        url_validation = self.validate_url(inputs)
        if url_validation:
            return url_validation

        status_code_validation = self.validate_status_code(inputs)
        if status_code_validation:
            return status_code_validation

        return None

    def assert_url_reachable(self, url: str, expected_status_code: int):
        try:
            res = request.urlopen(url)

            if res.getcode() != expected_status_code:
                return (
                    f"Unexpected status code for {url}: {res.getcode()}. "
                    f"Expected {expected_status_code}."
                )
        except error.HTTPError as e:
            if e.code == expected_status_code:
                return None

            return (
                f"Unexpected status code for {url}: {e.code}. "
                f"Expected {expected_status_code}."
            )
        except error.URLError as e:
            return f"Failed to reach {url}: {e.reason}."
        except Exception as e:
            return f"Failed to reach {url}: {e}"

        return None

    def get_url_from_outputs(self, output_name: str, state: Dict[str, object]):
        if "values" not in state or "outputs" not in state["values"]:
            raise ValueError(
                (
                    "Could not find 'values' in state. "
                    "Make sure terraform apply has been "
                    "run successfully."
                )
            )

        outputs = state["values"]["outputs"]
        if output_name not in outputs:
            raise ValueError(
                (
                    f"Could not find {output_name} in outputs. "
                    "Make sure your terraform code has an "
                    f"output named {output_name}."
                )
            )

        return outputs[output_name]["value"]

    def run_assertion(
        self, inputs: Dict[str, object], state: Dict[str, object]
    ):
        print("Running run_assertion for URLReachable")

        plugin_inputs = inputs
        url = plugin_inputs[self.FIELD_URL]
        from_outputs = plugin_inputs.get(self.FIELD_FROM_OUTPUTS, False)

        if from_outputs:
            url = self.get_url_from_outputs(url, state)

        expected_status_code = inputs.get(
            self.FIELD_STATUS_CODE, 200
        )

        return self.assert_url_reachable(url, expected_status_code)


def load_plugin():
    return URLReachableAssertionPlugin()

This registers both the "ExampleAssertion" and "URLReachable" assertions

[project]
name = "infra-tester-example-plugin"
version = "0.1.0"

[project.entry-points."infra_tester.assertion"]
ExampleAssertion = "example:load_plugin"
URLReachable = "url_reachable:load_plugin"

...
  - name: ExamplesCustomAssertions
    apply:
        assertions:
        - name: ThisOneAlwaysSucceeds
          type: ExampleAssertion

        - name: URLShouldBeReachable
          type: URLReachable
          url: https://www.schrodinger.com

        - name: URLRedirectsAreFollowedToReachableURL
          type: URLReachable
          url: http://www.schrodinger.com
          status_code: 200

        - name: OutputURLShouldBeReachable
          type: URLReachable
          url: sample_url
          from_outputs: true

        - name: URLShouldNotBeReachable
          type: URLReachable
          url: https://google.com/doesnotexist
          status_code: 404
...

. # example/ folder in the infra-tester git repo
└── plugin-example
    ├── pyproject.toml
    ├── setup.py
    └── src
        ├── example
        │   └── __init__.py
        └── url_reachable
            └── __init__.py

You can install the above package by running pip install -e <path/to/package>.

Testing Your Plugin

You may use the infra-tester-run-plugin command to execute your plugin independently to test its working. We also recommend writing unit tests for your plugins to catch any issues early on.

Run infra-tester-run-plugin -h to see how to use this CLI command.

Listing Available Plugins

You can list the available packages using the infra-tester-plugin-manager CLI command provided by the infra-tester-plugins package.

$ infra-tester-plugin-manager --list
ExampleAssertion
URLReachable

Conclusion

The plugin system provides a mechanism in which developers can write plugins in Python and then publish them so that other people who may not be familiar with Python or programming in general could still write tests in a simple YAML config.

As with all plugin systems which allows running external code, make sure you trust the plugin by auditing its code to ensure that unwanted code is not executed in your environment.

If you feel a plugin might be a better fit as an inbuilt assertion, please reach out by raising an issue in our GitHub repository.