From 4aec9d302fc86ac53433cf7564cc7447e49aadc6 Mon Sep 17 00:00:00 2001
From: "Steven B." <51370195+sdb9696@users.noreply.github.com>
Date: Tue, 29 Oct 2024 09:30:30 +0000
Subject: [PATCH] Allow enabling experimental devices from environment variable
 (#1194)

---
 devtools/dump_devinfo.py      |  4 ++--
 kasa/cli/main.py              | 14 ++++++++------
 kasa/device_factory.py        |  4 ++--
 kasa/experimental/__init__.py | 27 ++++++++++++++++++++++++++
 kasa/experimental/enabled.py  | 12 ------------
 kasa/tests/test_cli.py        | 36 +++++++++++++++++++++++++++++++++++
 6 files changed, 75 insertions(+), 22 deletions(-)
 delete mode 100644 kasa/experimental/enabled.py

diff --git a/devtools/dump_devinfo.py b/devtools/dump_devinfo.py
index f3b7810e..47d48454 100644
--- 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
index b721e984..a386fe4b 100755
--- a/kasa/cli/main.py
+++ b/kasa/cli/main.py
@@ -16,6 +16,7 @@ if TYPE_CHECKING:
     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
index 53ae1eff..d7b77843 100755
--- 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
index 60462246..388c5736 100644
--- 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
deleted file mode 100644
index 7679f97c..00000000
--- a/kasa/experimental/enabled.py
+++ /dev/null
@@ -1,12 +0,0 @@
-"""Package for experimental enabled."""
-
-
-class Enabled:
-    """Class for enabling experimental functionality."""
-
-    value = False
-
-    @classmethod
-    def set(cls, value):
-        """Set the enabled value."""
-        cls.value = value
diff --git a/kasa/tests/test_cli.py b/kasa/tests/test_cli.py
index bd93d430..80b5daaf 100644
--- 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