dify/scripts/stress-test/common/config_helper.py

#!/usr/bin/env python3

import json
from pathlib import Path
from typing import NotRequired, TypedDict


class AdminConfig(TypedDict):
    """Configuration for admin section."""

    username: str
    password: str
    base_url: str


class AuthConfig(TypedDict):
    """Configuration for authentication section."""

    access_token: str
    refresh_token: NotRequired[str]
    expires_at: NotRequired[int]


class AppConfig(TypedDict):
    """Configuration for app section."""

    app_id: str
    app_name: NotRequired[str]
    description: NotRequired[str]


class ApiKeyConfig(TypedDict):
    """Configuration for API key section."""

    token: str
    key_name: NotRequired[str]
    expires_at: NotRequired[int]


class StressTestState(TypedDict):
    """Complete stress test state structure."""

    admin: NotRequired[AdminConfig]
    auth: NotRequired[AuthConfig]
    app: NotRequired[AppConfig]
    api_key: NotRequired[ApiKeyConfig]


class ConfigHelper:
    _LEGACY_SECTION_MAP = {
        "admin_config": "admin",
        "token_config": "auth",
        "app_config": "app",
        "api_key_config": "api_key",
    }

    """Helper class for reading and writing configuration files."""

    def __init__(self, base_dir: Path | None = None):
        """Initialize ConfigHelper with base directory.

        Args:
            base_dir: Base directory for config files. If None, uses setup/config
        """
        if base_dir is None:
            # Default to config directory in setup folder
            base_dir = Path(__file__).parent.parent / "setup" / "config"
        self.base_dir = base_dir
        self.state_file = "stress_test_state.json"

    def ensure_config_dir(self) -> None:
        """Ensure the config directory exists."""
        self.base_dir.mkdir(exist_ok=True, parents=True)

    def get_config_path(self, filename: str) -> Path:
        """Get the full path for a config file.

        Args:
            filename: Name of the config file (e.g., 'admin_config.json')

        Returns:
            Full path to the config file
        """
        if not filename.endswith(".json"):
            filename += ".json"
        return self.base_dir / filename

    def read_config[T](self, filename: str) -> T | None:
        """Read a configuration file with generic return type.

        DEPRECATED: Use read_state() or get_state_section() for new code.
        This method provides backward compatibility.

        Args:
            filename: Name of the config file to read

        Returns:
            Configuration data of type T, or None if file doesn't exist
        """
        # Provide backward compatibility for old config names
        if filename in self._LEGACY_SECTION_MAP:
            section_data = self.get_state_section(self._LEGACY_SECTION_MAP[filename])
            return section_data  # type: ignore

        config_path = self.get_config_path(filename)

        if not config_path.exists():
            return None

        try:
            with open(config_path) as f:
                return json.load(f)  # type: ignore
        except (OSError, json.JSONDecodeError) as e:
            print(f"❌ Error reading {filename}: {e}")
            return None

    def write_config[T](self, filename: str, data: T) -> bool:
        """Write data to a configuration file.

        DEPRECATED: Use write_state() or update_state_section() for new code.
        This method provides backward compatibility.

        Args:
            filename: Name of the config file to write
            data: Data to save (must be JSON serializable)

        Returns:
            True if successful, False otherwise
        """
        # Provide backward compatibility for old config names
        if filename in self._LEGACY_SECTION_MAP:
            return self.update_state_section(
                self._LEGACY_SECTION_MAP[filename],
                data,  # type: ignore
            )

        self.ensure_config_dir()
        config_path = self.get_config_path(filename)

        try:
            with open(config_path, "w") as f:
                json.dump(data, f, indent=2)  # type: ignore
            return True
        except OSError as e:
            print(f"❌ Error writing {filename}: {e}")
            return False

    def config_exists(self, filename: str) -> bool:
        """Check if a config file exists.

        Args:
            filename: Name of the config file to check

        Returns:
            True if file exists, False otherwise
        """
        return self.get_config_path(filename).exists()

    def delete_config(self, filename: str) -> bool:
        """Delete a configuration file.

        Args:
            filename: Name of the config file to delete

        Returns:
            True if successful, False otherwise
        """
        config_path = self.get_config_path(filename)

        if not config_path.exists():
            return True  # Already doesn't exist

        try:
            config_path.unlink()
            return True
        except OSError as e:
            print(f"❌ Error deleting {filename}: {e}")
            return False

    def read_state(self) -> StressTestState | None:
        """Read the entire stress test state.

        Returns:
            Dictionary containing all state data, or None if file doesn't exist
        """
        state_path = self.get_config_path(self.state_file)
        if not state_path.exists():
            return None

        try:
            with open(state_path) as f:
                data = json.load(f)
                # Validate basic structure
                if not isinstance(data, dict):
                    print(f"❌ Invalid state format in {self.state_file}")
                    return None
                return data  # type: ignore
        except (OSError, json.JSONDecodeError) as e:
            print(f"❌ Error reading {self.state_file}: {e}")
            return None

    def write_state(self, data: StressTestState) -> bool:
        """Write the entire stress test state.

        Args:
            data: Dictionary containing all state data to save

        Returns:
            True if successful, False otherwise
        """
        self.ensure_config_dir()
        state_path = self.get_config_path(self.state_file)

        try:
            with open(state_path, "w") as f:
                json.dump(data, f, indent=2)
            return True
        except OSError as e:
            print(f"❌ Error writing {self.state_file}: {e}")
            return False

    def update_state_section[T](self, section: str, data: T) -> bool:
        """Update a specific section of the stress test state.

        Args:
            section: Name of the section to update (e.g., 'admin', 'auth', 'app', 'api_key')
            data: Section data to save

        Returns:
            True if successful, False otherwise
        """
        state = self.read_state() or {}
        state[section] = data  # type: ignore
        return self.write_state(state)  # type: ignore

    def get_state_section[T](self, section: str) -> T | None:
        """Get a specific section from the stress test state.

        Args:
            section: Name of the section to get (e.g., 'admin', 'auth', 'app', 'api_key')

        Returns:
            Section data of type T, or None if not found
        """
        state = self.read_state()
        if state:
            return state.get(section)  # type: ignore
        return None

    def get_token(self) -> str | None:
        """Get the access token from auth section.

        Returns:
            Access token string or None if not found
        """
        auth = self.get_state_section[AuthConfig]("auth")
        if auth:
            return auth.get("access_token")
        return None

    def get_app_id(self) -> str | None:
        """Get the app ID from app section.

        Returns:
            App ID string or None if not found
        """
        app = self.get_state_section[AppConfig]("app")
        if app:
            return app.get("app_id")
        return None

    def get_api_key(self) -> str | None:
        """Get the API key token from api_key section.

        Returns:
            API key token string or None if not found
        """
        api_key = self.get_state_section[ApiKeyConfig]("api_key")
        if api_key:
            return api_key.get("token")
        return None


# Create a default instance for convenience
config_helper = ConfigHelper()