"""Register test handlers."""
from __future__ import annotations
import logging
from typing import Dict, Type, Union, cast
from ..utils import load_object_from_string
from .handlers.base import LookupHandler
from .handlers.cfn import CfnLookup
from .handlers.ecr import EcrLookup
from .handlers.env import EnvLookup
from .handlers.random_string import RandomStringLookup
from .handlers.ssm import SsmLookup
from .handlers.var import VarLookup
RUNWAY_LOOKUP_HANDLERS: Dict[str, Type[LookupHandler]] = {}
LOGGER = logging.getLogger(__name__)
[docs]def register_lookup_handler(
lookup_type: str, handler_or_path: Union[str, Type[LookupHandler]]
) -> None:
"""Register a lookup handler.
Args:
lookup_type: Name to register the handler under
handler_or_path: a function or a path to a handler
"""
handler = handler_or_path
if isinstance(handler_or_path, str):
handler = cast(type, load_object_from_string(handler_or_path))
else:
handler = handler_or_path
try:
if issubclass(handler, LookupHandler):
RUNWAY_LOOKUP_HANDLERS[lookup_type] = handler
return
except Exception: # pylint: disable=broad-except
LOGGER.debug("failed to validate lookup handler", exc_info=True)
raise TypeError(
f"lookup {handler_or_path} must be a subclass of "
"runway.lookups.handlers.base.LookupHandler"
)
[docs]def unregister_lookup_handler(lookup_type: str) -> None:
"""Unregister the specified test type.
This is useful when testing various lookup types if you want to unregister
the lookup type after the test runs.
Args:
lookup_type: Name of the lookup type to unregister
"""
RUNWAY_LOOKUP_HANDLERS.pop(lookup_type, None)
register_lookup_handler(CfnLookup.TYPE_NAME, CfnLookup)
register_lookup_handler(EcrLookup.TYPE_NAME, EcrLookup)
register_lookup_handler(EnvLookup.TYPE_NAME, EnvLookup)
register_lookup_handler(RandomStringLookup.TYPE_NAME, RandomStringLookup)
register_lookup_handler(SsmLookup.TYPE_NAME, SsmLookup)
register_lookup_handler(VarLookup.TYPE_NAME, VarLookup)