diff --git a/chaoslib/control/__init__.py b/chaoslib/control/__init__.py index ca34fb9..a8057eb 100644 --- a/chaoslib/control/__init__.py +++ b/chaoslib/control/__init__.py @@ -3,15 +3,17 @@ from copy import deepcopy from typing import List, Union +import contextvars from logzero import logger from chaoslib.control.python import apply_python_control, cleanup_control, \ initialize_control, validate_python_control from chaoslib.exceptions import InterruptExecution, InvalidControl +from chaoslib.settings import get_loaded_settings from chaoslib.types import Settings -from chaoslib.types import Activity, Configuration, \ +from chaoslib.types import Activity, Configuration, Control as ControlType, \ Experiment, Hypothesis, Journal, Run, Secrets -import contextvars + __all__ = ["controls", "initialize_controls", "cleanup_controls", "validate_controls", "Control", "initialize_global_controls", @@ -114,21 +116,16 @@ def initialize_global_controls(settings: Settings): """ Load and initialize controls declared in the settings """ - auths = settings.get('auths', []) controls = [] for name, control in settings.get("controls", {}).items(): control['name'] = name logger.debug("Initializing global control '{}'".format(name)) - auth_secrets = { - auth: deepcopy(auths[auth]) - for auth in control.get( - "secrets", {}).get('auths', []) if auth in auths - } provider = control.get("provider") if provider and provider["type"] == "python": initialize_control( - control, configuration=None, secrets=auth_secrets) + control, configuration=None, secrets=None, + settings=settings) controls.append(control) global_controls.set(controls) @@ -138,6 +135,8 @@ def cleanup_global_controls(): Unload and cleanup global controls """ controls = global_controls.get() + global_controls.set([]) + for control in controls: name = control['name'] logger.debug("Cleaning up global control '{}'".format(name)) @@ -147,6 +146,13 @@ def cleanup_global_controls(): cleanup_control(control) +def get_global_controls() -> List[ControlType]: + """ + All the controls loaded from the settings. + """ + return global_controls.get() + + class Control: def begin(self, level: str, experiment: Experiment, context: Union[Activity, Hypothesis, Experiment], @@ -273,6 +279,7 @@ def apply_controls(level: str, experiment: Experiment, the `"activity"` when it must be an activity. The `scope` is one of `"before", "after"` and the `state` is only set on `"after"` scope. """ + settings = get_loaded_settings() or None controls = get_context_controls(level, experiment, context) if not controls: return @@ -295,7 +302,8 @@ def apply_controls(level: str, experiment: Experiment, apply_python_control( level="{}-{}".format(level, scope), control=control, context=context, state=state, experiment=experiment, - configuration=configuration, secrets=secrets) + configuration=configuration, secrets=secrets, + settings=settings) except InterruptExecution: logger.debug( "{}-control '{}' interrupted the execution".format( diff --git a/chaoslib/control/python.py b/chaoslib/control/python.py index 26e4a4c..0236fc4 100644 --- a/chaoslib/control/python.py +++ b/chaoslib/control/python.py @@ -9,7 +9,7 @@ from chaoslib import substitute from chaoslib.exceptions import InvalidActivity from chaoslib.types import Activity, Configuration, Control, Experiment, \ - Journal, Run, Secrets + Journal, Run, Secrets, Settings __all__ = ["apply_python_control", "cleanup_control", "initialize_control", @@ -29,14 +29,27 @@ def initialize_control(control: Control, configuration: Configuration, - secrets: Secrets): + secrets: Secrets, settings: Settings = None): """ Initialize a control by calling its `configure_control` function. """ func = load_func(control, "configure_control") if not func: return - func(configuration, secrets) + + arguments = {} + sig = inspect.signature(func) + + if "secrets" in sig.parameters: + arguments["secrets"] = secrets + + if "config" in sig.parameters: + arguments["config"] = configuration + + if "settings" in sig.parameters: + arguments["settings"] = settings + + func(**arguments) def cleanup_control(control: Control): @@ -73,7 +86,7 @@ def apply_python_control(level: str, control: Control, experiment: Experiment, context: Union[Activity, Experiment], state: Union[Journal, Run, List[Run]] = None, configuration: Configuration = None, - secrets: Secrets = None): + secrets: Secrets = None, settings: Settings = None): """ Apply a control by calling a function matching the given level. """ @@ -103,6 +116,12 @@ def apply_python_control(level: str, control: Control, experiment: Experiment, if "experiment" in sig.parameters: arguments["experiment"] = experiment + if "extensions" in sig.parameters: + arguments["extensions"] = experiment.get("extensions") + + if "settings" in sig.parameters: + arguments["settings"] = settings + func(context=context, **arguments) diff --git a/tests/fixtures/controls/dummy.py b/tests/fixtures/controls/dummy.py index f2b39bc..9bb718e 100644 --- a/tests/fixtures/controls/dummy.py +++ b/tests/fixtures/controls/dummy.py @@ -7,10 +7,13 @@ value_from_config = None -def configure_control(config: Configuration, secrets = Secrets): +def configure_control(config: Configuration, secrets: Secrets, + settings: Settings): global value_from_config - print(config) - value_from_config = config.get("dummy-key", "default") + if config: + value_from_config = config.get("dummy-key", "default") + elif settings: + value_from_config = settings.get("dummy-key", "default") def cleanup_control(): @@ -32,7 +35,7 @@ def before_hypothesis_control(context: Hypothesis, **kwargs): def after_hypothesis_control(context: Hypothesis, - state: Dict[str, Any], **kwargs): + state: Dict[str, Any], **kwargs): context["after_hypothesis_control"] = True state["after_hypothesis_control"] = True diff --git a/tests/fixtures/controls/dummy_with_experiment.py b/tests/fixtures/controls/dummy_with_experiment.py index 17a6d61..67594f0 100644 --- a/tests/fixtures/controls/dummy_with_experiment.py +++ b/tests/fixtures/controls/dummy_with_experiment.py @@ -7,9 +7,8 @@ value_from_config = None -def configure_control(config: Configuration, secrets = Secrets): +def configure_control(config: Configuration, secrets: Secrets): global value_from_config - print(config) value_from_config = config.get("dummy-key", "default") diff --git a/tests/fixtures/experiments.py b/tests/fixtures/experiments.py index 7667167..4a031ec 100644 --- a/tests/fixtures/experiments.py +++ b/tests/fixtures/experiments.py @@ -229,7 +229,7 @@ "steady-state-hypothesis": { "title": "hello", "probes": [ - deepcopy(PythonModuleProbeWithBoolTolerance) + deepcopy(PythonModuleProbeWithBoolTolerance) ] }, "method": [ diff --git a/tests/test_control.py b/tests/test_control.py index cb2ca95..62503f7 100644 --- a/tests/test_control.py +++ b/tests/test_control.py @@ -12,7 +12,8 @@ from chaoslib.activity import execute_activity from chaoslib.control import initialize_controls, cleanup_controls, \ - validate_controls, controls, get_all_activities, get_context_controls + validate_controls, controls, get_all_activities, get_context_controls, \ + initialize_global_controls, cleanup_global_controls, get_global_controls from chaoslib.control.python import validate_python_control from chaoslib.exceptions import InterruptExecution, InvalidActivity from chaoslib.experiment import run_experiment @@ -258,7 +259,6 @@ def test_controls_are_applied_at_various_levels(): run_experiment(exp) activities = get_all_activities(exp) for activity in activities: - print(activity) if "controls" in activity: assert activity["before_activity_control"] is True assert activity["after_activity_control"] is True @@ -274,3 +274,64 @@ def test_controls_are_applied_when_they_are_not_top_level(): if "controls" in activity: assert activity["before_activity_control"] is True assert activity["after_activity_control"] is True + + +def test_load_global_controls_from_settings(): + exp = deepcopy(experiments.ExperimentWithControls) + + try: + run_experiment(exp) + activities = get_all_activities(exp) + for activity in activities: + if "controls" in activity: + assert "before_activity_control" not in activity + assert "after_activity_control" not in activity + finally: + cleanup_global_controls() + + initialize_global_controls({ + "dummy-key": "blah", + "controls": { + "dummy": { + "provider": { + "type": "python", + "module": "fixtures.controls.dummy" + } + } + } + }) + + try: + run_experiment(exp) + activities = get_all_activities(exp) + for activity in activities: + if "controls" in activity: + assert activity["before_activity_control"] is True + assert activity["after_activity_control"] is True + finally: + cleanup_global_controls() + + +def test_get_globally_loaded_controls_from_settings(): + assert get_global_controls() == [] + + initialize_global_controls({ + "controls": { + "dummy": { + "provider": { + "type": "python", + "module": "fixtures.controls.dummy" + } + } + } + }) + + try: + ctrls = get_global_controls() + assert len(ctrls) == 1 + assert ctrls[0]["name"] == "dummy" + assert ctrls[0]["provider"]["type"] == "python" + assert ctrls[0]["provider"]["module"] == "fixtures.controls.dummy" + finally: + cleanup_global_controls() + assert get_global_controls() == []