test_client.py

#!/usr/bin/env python3
import json

import requests


class PolicyGateClient:
    def __init__(self, base_url: str = "http://localhost:8080"):
        self.base_url = base_url

    def switch_profile(self, profile_name: str) -> dict:
        """
        Switch to a different policy profile.

        Args:
            profile_name: Target profile name

        Returns:
            Response dictionary
        """
        response = requests.post(
            f"{self.base_url}/switch-profile",
            json={"profile": profile_name},
            timeout=5,
        )
        return response.json()

    def enforce_request(self, request_data: dict) -> dict:
        """
        Submit a request for policy enforcement.

        Args:
            request_data: Request to validate

        Returns:
            Enforcement results
        """
        response = requests.post(
            f"{self.base_url}/enforce",
            json=request_data,
            timeout=5,
        )
        return response.json()

    def get_status(self) -> dict:
        """Get current policy gate status."""
        response = requests.get(f"{self.base_url}/status", timeout=5)
        return response.json()


# Create test scenarios for each industry profile
def test_healthcare_profile():
    """Test healthcare profile enforcement."""
    client = PolicyGateClient()

    print("\n=== Testing Healthcare Profile ===")
    print(json.dumps(client.switch_profile("healthcare"), indent=2))

    non_compliant = {
        "data_type": "PHI",
        "encrypted": False,
        "access_logged": True,
        "retention_days": 2000,
    }
    print(json.dumps(client.enforce_request(non_compliant), indent=2))

    compliant = {
        "data_type": "PHI",
        "encrypted": True,
        "access_logged": True,
        "retention_days": 2000,
    }
    print(json.dumps(client.enforce_request(compliant), indent=2))


def test_finance_profile():
    """Test finance profile enforcement."""
    client = PolicyGateClient()

    print("\n=== Testing Finance Profile ===")
    print(json.dumps(client.switch_profile("finance"), indent=2))

    non_compliant = {
        "data_type": "cardholder_data",
        "encrypted": True,
        "network_segmented": False,
        "password_length": 8,
    }
    print(json.dumps(client.enforce_request(non_compliant), indent=2))

    compliant = {
        "data_type": "cardholder_data",
        "encrypted": True,
        "network_segmented": True,
        "password_length": 14,
    }
    print(json.dumps(client.enforce_request(compliant), indent=2))


def test_retail_profile():
    """Test retail profile enforcement."""
    client = PolicyGateClient()

    print("\n=== Testing Retail Profile ===")
    print(json.dumps(client.switch_profile("retail"), indent=2))

    non_compliant = {
        "data_type": "customer_pii",
        "encrypted": True,
        "current_requests": 1500,
    }
    print(json.dumps(client.enforce_request(non_compliant), indent=2))

    compliant = {
        "data_type": "customer_pii",
        "encrypted": True,
        "current_requests": 500,
    }
    print(json.dumps(client.enforce_request(compliant), indent=2))


if __name__ == "__main__":
    test_healthcare_profile()
    test_finance_profile()
    test_retail_profile()