diff --git a/src/xpk/parser/cluster.py b/src/xpk/parser/cluster.py index 97f6b0408..a4e416769 100644 --- a/src/xpk/parser/cluster.py +++ b/src/xpk/parser/cluster.py @@ -31,7 +31,7 @@ from ..core.vertex import DEFAULT_VERTEX_TENSORBOARD_NAME from .common import add_shared_arguments, ParserOrArgumentGroup from .validators import name_type -from ..utils.feature_flags import SUB_SLICING_ENABLED +from ..utils.feature_flags import FeatureFlags def set_cluster_parser(cluster_parser: ArgumentParser): @@ -143,7 +143,7 @@ def set_cluster_create_parser(cluster_create_parser: ArgumentParser): ' enable cluster to accept Pathways workloads.' ), ) - if SUB_SLICING_ENABLED: + if FeatureFlags.SUB_SLICING_ENABLED: cluster_create_optional_arguments.add_argument( '--sub-slicing', action='store_true', diff --git a/src/xpk/parser/cluster_test.py b/src/xpk/parser/cluster_test.py index 86ea15f48..ec1d81b6b 100644 --- a/src/xpk/parser/cluster_test.py +++ b/src/xpk/parser/cluster_test.py @@ -17,15 +17,16 @@ import argparse from xpk.parser.cluster import set_cluster_create_parser import pytest +from ..utils.feature_flags import FeatureFlags @pytest.fixture(autouse=True) -def with_sub_slicing_enabled(mocker): - mocker.patch("xpk.parser.cluster.SUB_SLICING_ENABLED", True) +def with_sub_slicing_enabled(): + FeatureFlags.SUB_SLICING_ENABLED = True -def test_cluster_create_sub_slicing_is_hidden_with_flag_off(mocker): - mocker.patch("xpk.parser.cluster.SUB_SLICING_ENABLED", False) +def test_cluster_create_sub_slicing_is_hidden_with_flag_off(): + FeatureFlags.SUB_SLICING_ENABLED = False parser = argparse.ArgumentParser() set_cluster_create_parser(parser) diff --git a/src/xpk/utils/feature_flags.py b/src/xpk/utils/feature_flags.py index 089394571..6c6d22fff 100644 --- a/src/xpk/utils/feature_flags.py +++ b/src/xpk/utils/feature_flags.py @@ -17,8 +17,12 @@ import os -def __getBooleanFlag(flag: str, default: bool) -> bool: +def _get_boolean_flag(flag: str, default: bool) -> bool: return os.getenv(flag, str(default)).lower() == "true" -SUB_SLICING_ENABLED = __getBooleanFlag("SUB_SLICING_ENABLED", default=False) +class _FeatureFlags: + SUB_SLICING_ENABLED = _get_boolean_flag("SUB_SLICING_ENABLED", default=False) + + +FeatureFlags = _FeatureFlags()