Allow feature toggling from environment variables

add possibility for optional toggles

src/ert/__main__.py

@@ -6,8 +6,7 @@
 import re
 import sys
 from argparse import ArgumentParser, ArgumentTypeError
-from contextlib import contextmanager
-from typing import Any, Dict, Generator, Optional, Sequence, Union
+from typing import Any, Dict, Optional, Sequence, Union
 from uuid import UUID
 
 import yaml
@@ -494,16 +493,6 @@ def ert_parser(parser: Optional[ArgumentParser], args: Sequence[str]) -> Namespa
     )
 
 
-@contextmanager
-def start_ert_server(mode: str) -> Generator[None, None, None]:
-    if mode in ("api", "vis") or not FeatureToggling.is_enabled("new-storage"):
-        yield
-        return
-
-    with StorageService.start_server():
-        yield
-
-
 def log_process_usage() -> None:
     try:
         import resource
@@ -571,7 +560,7 @@ def main() -> None:
 
     FeatureToggling.update_from_args(args)
     try:
-        with start_ert_server(args.mode), ErtPluginContext() as context:
+        with ErtPluginContext() as context:
             context.plugin_manager.add_logging_handle_to_root(logging.getLogger())
             logger.info(f"Running ert with {args}")
             args.func(args, context.plugin_manager)

src/ert/shared/feature_toggling.py

@@ -1,71 +1,117 @@
 import logging
+import os
 from argparse import ArgumentParser
 from copy import deepcopy
-from typing import Optional
+from typing import Dict, Optional, Union
 
 from ert.namespace import Namespace
 
 
 class _Feature:
-    def __init__(self, default_enabled: bool, msg: Optional[str] = None) -> None:
-        self.is_enabled = default_enabled
+    def __init__(
+        self, default: Optional[bool], msg: Optional[str] = None, optional: bool = False
+    ) -> None:
+        self._value = default
         self.msg = msg
+        self.optional = optional
+
+    def validate_value(self, value: Union[bool, str, None]) -> Optional[bool]:
+        if type(value) is bool or value is None:
+            return value
+        elif value.lower() in ["true", "1"]:
+            return True
+        elif value.lower() in ["false", "0"]:
+            return False
+        elif self.optional and value.lower() in ["default", ""]:
+            return None
+        else:
+            raise ValueError(
+                f"This option can only be set to {'True/1, False/0 or Default/<empty>' if self.optional else 'True/1 or False/0'}"
+            )
+
+    @property
+    def value(self) -> Optional[bool]:
+        return self._value
+
+    @value.setter
+    def value(self, value: Optional[bool]) -> None:
+        self._value = self.validate_value(value)
 
 
 class FeatureToggling:
-    _conf_original = {
-        "new-storage": _Feature(
-            default_enabled=False,
-            msg=(
-                "The new storage solution is experimental! "
-                "Thank you for testing our new features."
-            ),
-        ),
+    _conf_original: Dict[str, _Feature] = {
         "scheduler": _Feature(
-            default_enabled=False,
-            msg="Use Scheduler instead of JobQueue",
+            default=None,
+            msg="Default value for use of Scheduler has been overridden\n"
+            "This is experimental and may cause problems",
+            optional=True,
         ),
     }
 
     _conf = deepcopy(_conf_original)
 
     @staticmethod
     def is_enabled(feature_name: str) -> bool:
-        return FeatureToggling._conf[feature_name].is_enabled
+        return FeatureToggling._conf[feature_name].value is True
+
+    @staticmethod
+    def value(feature_name: str) -> Optional[bool]:
+        return FeatureToggling._conf[feature_name].value
 
     @staticmethod
     def add_feature_toggling_args(parser: ArgumentParser) -> None:
-        for feature_name in FeatureToggling._conf:
-            parser.add_argument(
-                f"--{FeatureToggling._get_arg_name(feature_name)}",
-                action="store_true",
-                help=f"Toggle {feature_name} (Warning: This is experimental)",
-                default=False,
-            )
+        for name, feature in FeatureToggling._conf.items():
+            env_var_name = f"ERT_FEATURE_{name.replace('-', '_').upper()}"
+            env_value: Union[bool, str, None] = None
+            if env_var_name in os.environ:
+                try:
+                    feature.value = feature.validate_value(os.environ[env_var_name])
+                except ValueError as e:
+                    # TODO: this is a bit spammy. It will get called 6 times for each incorrect env var.
+                    logging.getLogger().warning(
+                        f"Failed to set {env_var_name} to '{os.environ[env_var_name]}'. {e}"
+                    )
+
+            if not feature.optional:
+                parser.add_argument(
+                    f"--{'disable' if feature.value else 'enable'}-{name}",
+                    action="store_false" if feature.value else "store_true",
+                    help=f"Toggle {name} (Warning: This is experimental)",
+                    dest=f"feature-{name}",
+                    default=env_value if env_value is not None else feature.value,
+                )
+            else:
+                group = parser.add_mutually_exclusive_group()
+                group.add_argument(
+                    f"--enable-{name}",
+                    action="store_true",
+                    help=f"Enable {name}",
+                    dest=f"feature-{name}",
+                    default=feature.value,
+                )
+                group.add_argument(
+                    f"--disable-{name}",
+                    action="store_false",
+                    help=f"Disable {name}",
+                    dest=f"feature-{name}",
+                    default=feature.value,
+                )
 
     @staticmethod
     def update_from_args(args: Namespace) -> None:
-        args_dict = vars(args)
-        for feature_name in FeatureToggling._conf:
-            arg_name = FeatureToggling._get_arg_name(feature_name)
-            feature_name_escaped = arg_name.replace("-", "_")
-
-            if feature_name_escaped in args_dict and args_dict[feature_name_escaped]:
-                current_state = FeatureToggling._conf[feature_name].is_enabled
-                FeatureToggling._conf[feature_name].is_enabled = not current_state
-
-            if (
-                FeatureToggling._conf[feature_name].is_enabled
-                and FeatureToggling._conf[feature_name].msg is not None
-            ):
-                logger = logging.getLogger()
-                logger.warning(FeatureToggling._conf[feature_name].msg)
+        pattern = "feature-"
+        feature_args = [arg for arg in vars(args).items() if arg[0].startswith(pattern)]
+        for name, value in feature_args:
+            name = name[len(pattern) :]
+            if name in FeatureToggling._conf:
+                FeatureToggling._conf[name].value = value
 
-    @staticmethod
-    def _get_arg_name(feature_name: str) -> str:
-        default_state = FeatureToggling._conf[feature_name].is_enabled
-        arg_default_state = "disable" if default_state else "enable"
-        return f"{arg_default_state}-{feature_name}"
+        # Print warnings for enabled features.
+        for name, feature in FeatureToggling._conf.items():
+            if FeatureToggling.is_enabled(name) and feature.msg is not None:
+                logging.getLogger().warning(
+                    f"{feature.msg}\nValue is set to {feature.value}"
+                )
 
     @staticmethod
     def reset() -> None:

tests/conftest.py

@@ -265,7 +265,7 @@ def try_queue_and_scheduler(request, monkeypatch):
         _ = get_event_loop()
 
     monkeypatch.setattr(
-        FeatureToggling._conf["scheduler"], "is_enabled", should_enable_scheduler
+        FeatureToggling._conf["scheduler"], "_value", should_enable_scheduler
     )
     yield
     monkeypatch.undo()

tests/unit_tests/shared/test_feature_toggling.py

@@ -0,0 +1,125 @@
+import os
+from unittest.mock import Mock
+
+import pytest
+
+import ert.__main__
+from ert.__main__ import ert_parser
+from ert.cli import TEST_RUN_MODE
+from ert.shared.feature_toggling import FeatureToggling
+
+
+def is_default(feature_name: str) -> bool:
+    return (
+        FeatureToggling._conf[feature_name].value
+        == FeatureToggling._conf_original[feature_name].value
+    )
+
+
+def feature_to_env_name(feature_name: str) -> str:
+    return f"ERT_FEATURE_{feature_name.replace('-', '_').upper()}"
+
+
+@pytest.fixture(autouse=True)
+def mocked_valid_file(monkeypatch):
+    monkeypatch.setattr(
+        ert.__main__, "valid_file", Mock(return_value="not_a_real_config.ert")
+    )
+
+
+@pytest.fixture(autouse=True)
+def reset_feature_toggling():
+    FeatureToggling.reset()
+
+
+@pytest.fixture(autouse=True)
+def reset_environment_vars(monkeypatch):
+    monkeypatch.setattr(os, "environ", {})
+
+
+def test_feature_toggling_from_env():
+    os.environ["ERT_FEATURE_SCHEDULER"] = "True"
+
+    parsed = ert_parser(
+        None,
+        [
+            TEST_RUN_MODE,
+            "not_a_real_config.ert",
+        ],
+    )
+    FeatureToggling.update_from_args(parsed)
+
+    assert FeatureToggling.is_enabled("scheduler")
+
+
+def test_feature_toggling_from_args():
+    parsed = ert_parser(
+        None,
+        [
+            TEST_RUN_MODE,
+            "--disable-scheduler",
+            "not_a_real_config.ert",
+        ],
+    )
+    FeatureToggling.update_from_args(parsed)
+
+    assert not FeatureToggling.is_enabled("scheduler")
+
+
+@pytest.mark.parametrize(
+    "environment_vars, arguments, expected",
+    [
+        (
+            {"ERT_FEATURE_SCHEDULER": "False"},
+            ["--enable-scheduler"],
+            {"scheduler": True},
+        ),
+        (
+            {"ERT_FEATURE_SCHEDULER": "True"},
+            ["--disable-scheduler"],
+            {"scheduler": False},
+        ),
+        (
+            {"ERT_FEATURE_SCHEDULER": ""},
+            ["--disable-scheduler"],
+            {"scheduler": False},
+        ),
+    ],
+)
+def test_feature_toggling_both(environment_vars, arguments, expected):
+    for key, value in environment_vars.items():
+        os.environ[key] = value
+
+    parsed = ert_parser(
+        None,
+        [
+            TEST_RUN_MODE,
+            *arguments,
+            "not_a_real_config.ert",
+        ],
+    )
+
+    FeatureToggling.update_from_args(parsed)
+
+    for key, value in expected.items():
+        assert FeatureToggling.value(key) == value
+
+
+@pytest.mark.parametrize(
+    "feature_name, value",
+    [
+        ("scheduler", "incorrect"),
+    ],
+)
+def test_feature_toggling_incorrect_input(feature_name, value):
+    os.environ[feature_to_env_name(feature_name)] = value
+    parsed = ert_parser(
+        None,
+        [
+            TEST_RUN_MODE,
+            "not_a_real_config.ert",
+        ],
+    )
+
+    FeatureToggling.update_from_args(parsed)
+    assert is_default(feature_name)

tests/unit_tests/shared/test_main_entry.py

@@ -13,7 +13,6 @@ def test_main_logging(monkeypatch, caplog):
     parser_mock.func.side_effect = ValueError("This is a test")
     monkeypatch.setattr(logging.config, "dictConfig", MagicMock())
     monkeypatch.setattr(main, "ert_parser", MagicMock(return_value=parser_mock))
-    monkeypatch.setattr(main, "start_ert_server", MagicMock())
     monkeypatch.setattr(main, "ErtPluginContext", MagicMock())
     monkeypatch.setattr(sys, "argv", ["ert", "test_run", "config.ert"])
     with pytest.raises(
@@ -28,7 +27,6 @@ def test_main_logging_argparse(monkeypatch, caplog):
     monkeypatch.setattr(logging.config, "dictConfig", MagicMock())
     monkeypatch.setattr(main, "valid_file", MagicMock(return_value=True))
     monkeypatch.setattr(main, "run_cli", MagicMock())
-    monkeypatch.setattr(main, "start_ert_server", MagicMock())
     monkeypatch.setattr(main, "ErtPluginContext", MagicMock())
     monkeypatch.setattr(sys, "argv", ["ert", "test_run", "config.ert"])
     with caplog.at_level(logging.INFO):