httphound/example_tests/02-basic_header_check.py

from pathlib import Path
from proxytester.main import BaseProxyTest, BackendConfig, ProxyConfig


class FailHeaderCheck(BaseProxyTest):
    """Minimal test to demonstrate the basic header checks functionality

    This test should fail as the backend or the reverse proxy aren't configured
    to answer this header.
    """
    def __init__(self):
        super().__init__()
        self.description = "Client custom header (Fail)"
        self.url = "http://localhost:4242/"

        # expected_headers must match with the headers received by
        # the client or will fail
        self.expected_headers = {
            "x-test": "1234",
        }

        self.proxy_config = ProxyConfig(
            binary_path=Path.home() / "bin/haproxy"
        )

    async def run_test(self):
        await self.make_request()
        return True


class PassHeaderCheck(BaseProxyTest):
    """Minimal test to demonstrate the basic header checks functionality

    This test should succeed.
    """
    def __init__(self):
        super().__init__()
        self.description = "Client custom header (Pass)"
        self.url = "http://localhost:4242/"

        # expected_headers must match with the headers received by
        # the client or will fail
        self.expected_headers = {
            "x-test": "1234",
        }

        self.proxy_config = ProxyConfig(
            binary_path=Path.home() / "bin/haproxy"
        )

        # Defining a backend configuration with headers
        # to answer to the client
        self.backend_config = BackendConfig(
            response_headers={"X-Test": "1234"},
        )

    async def run_test(self):
        """The run_test() method is always the same as all logic
        is defined entirely in the classes configuration.
        """
        await self.make_request()
        return True