python-kasa · sdb9696 · Oct 29, 2024 · Oct 25, 2024 · Oct 25, 2024 · Oct 25, 2024
diff --git a/devtools/dump_devinfo.py b/devtools/dump_devinfo.py
@@ -309,9 +309,9 @@ async def cli(
     if debug:
         logging.basicConfig(level=logging.DEBUG)
 
-    from kasa.experimental.enabled import Enabled
+    from kasa.experimental import Experimental
 
-    Enabled.set(True)
+    Experimental.set_enabled(True)
 
     credentials = Credentials(username=username, password=password)
     if host is not None:

diff --git a/kasa/cli/main.py b/kasa/cli/main.py
@@ -16,6 +16,7 @@
     from kasa import Device
 
 from kasa.deviceconfig import DeviceEncryptionType
+from kasa.experimental import Experimental
 
 from .common import (
     SKIP_UPDATE_COMMANDS,
@@ -220,11 +221,11 @@ def _legacy_type_to_class(_type):
     help="Hashed credentials used to authenticate to the device.",
 )
 @click.option(
-    "--experimental",
-    default=False,
+    "--experimental/--no-experimental",
+    default=None,
     is_flag=True,
     type=bool,
-    envvar="KASA_EXPERIMENTAL",
+    envvar=Experimental.ENV_VAR,
     help="Enable experimental mode for devices not yet fully supported.",
 )
 @click.version_option(package_name="python-kasa")
@@ -260,10 +261,11 @@ async def cli(
     if target != DEFAULT_TARGET and host:
         error("--target is not a valid option for single host discovery")
 
-    if experimental:
-        from kasa.experimental.enabled import Enabled
+    if experimental is not None:
+        Experimental.set_enabled(experimental)
 
-        Enabled.set(True)
+    if Experimental.enabled():
+        echo("Experimental support is enabled")
 
     logging_config: dict[str, Any] = {
         "level": logging.DEBUG if debug > 0 else logging.INFO

diff --git a/kasa/device_factory.py b/kasa/device_factory.py
@@ -214,9 +214,9 @@ def get_protocol(
         "SMART.KLAP": (SmartProtocol, KlapTransportV2),
     }
     if not (prot_tran_cls := supported_device_protocols.get(protocol_transport_key)):
-        from .experimental.enabled import Enabled
+        from .experimental import Experimental
 
-        if Enabled.value and protocol_transport_key == "SMART.AES.HTTPS":
+        if Experimental.enabled() and protocol_transport_key == "SMART.AES.HTTPS":
             prot_tran_cls = (SmartCameraProtocol, SslAesTransport)
         else:
             return None

diff --git a/kasa/experimental/__init__.py b/kasa/experimental/__init__.py
@@ -1 +1,28 @@
 """Package for experimental."""
+
+from __future__ import annotations
+
+import os
+
+
+class Experimental:
+    """Class for enabling experimental functionality."""
+
+    _enabled: bool | None = None
+    ENV_VAR = "KASA_EXPERIMENTAL"
+
+    @classmethod
+    def set_enabled(cls, enabled):
+        """Set the enabled value."""
+        cls._enabled = enabled
+
+    @classmethod
+    def enabled(cls):
+        """Get the enabled value."""
+        if cls._enabled is not None:
+            return cls._enabled
+
+        if env_var := os.getenv(cls.ENV_VAR):
+            return env_var.lower() in {"true", "1", "t", "on"}
+
+        return False
diff --git a/kasa/experimental/enabled.py b/kasa/experimental/enabled.py
diff --git a/kasa/tests/test_cli.py b/kasa/tests/test_cli.py
@@ -1232,3 +1232,39 @@ async def test_discover_config_invalid(mocker, runner):
     )
     assert res.exit_code == 1
     assert "--target is not a valid option for single host discovery" in res.output
+
+
+@pytest.mark.parametrize(
+    ("option", "env_var_value", "expectation"),
+    [
+        pytest.param("--experimental", None, True),
+        pytest.param("--experimental", "false", True),
+        pytest.param(None, None, False),
+        pytest.param(None, "true", True),
+        pytest.param(None, "false", False),
+        pytest.param("--no-experimental", "true", False),
+    ],
+)
+async def test_experimental_flags(mocker, option, env_var_value, expectation):
+    """Test the experimental flag is set correctly."""
+    mocker.patch("kasa.discover.Discover.try_connect_all", return_value=None)
+
+    # reset the class internal variable
+    from kasa.experimental import Experimental
+
+    Experimental._enabled = None
+
+    KASA_VARS = {k: None for k, v in os.environ.items() if k.startswith("KASA_")}
+    if env_var_value:
+        KASA_VARS["KASA_EXPERIMENTAL"] = env_var_value
+    args = [
+        "--host",
+        "127.0.0.2",
+        "discover",
+        "config",
+    ]
+    if option:
+        args.insert(0, option)
+    runner = CliRunner(env=KASA_VARS)
+    res = await runner.invoke(cli, args)
+    assert ("Experimental support is enabled" in res.output) is expectation