diff --git a/asusrouter/modules/led.py b/asusrouter/modules/led.py
index 4bab238..9ee855a 100644
--- a/asusrouter/modules/led.py
+++ b/asusrouter/modules/led.py
@@ -3,10 +3,9 @@
 from __future__ import annotations
 
 from enum import IntEnum
-from typing import Any, Awaitable, Callable, Optional
+from typing import Any, Awaitable, Callable
 
 from asusrouter.modules.endpoint import Endpoint
-from asusrouter.modules.identity import AsusDevice
 
 
 class AsusLED(IntEnum):
@@ -20,9 +19,7 @@ class AsusLED(IntEnum):
 async def set_state(
     callback: Callable[..., Awaitable[bool]],
     state: AsusLED,
-    arguments: Optional[dict[str, Any]] = None,
-    expect_modify: bool = False,
-    _: Optional[dict[Any, Any]] = None,
+    **kwargs: Any,
 ) -> bool:
     """Set the LED state."""
 
@@ -34,17 +31,19 @@ async def set_state(
         service="start_ctrl_led",
         arguments=arguments,
         apply=True,
-        expect_modify=expect_modify,
+        expect_modify=kwargs.get("expect_modify", False),
     )
 
 
 async def keep_state(
     callback: Callable[..., Awaitable[bool]],
     state: AsusLED = AsusLED.ON,
-    identity: Optional[AsusDevice] = None,
+    **kwargs: Any,
 ) -> bool:
     """Keep the LED state."""
 
+    identity = kwargs.get("identity")
+
     # Check if identity is available and if endpoints are defined
     if identity is None or not identity.endpoints:
         return False
@@ -57,7 +56,7 @@ async def keep_state(
         return False
 
     # Toggle the LED
-    await set_state(callback, AsusLED.ON)
-    await set_state(callback, AsusLED.OFF)
+    await set_state(callback, AsusLED.ON, **kwargs)
+    await set_state(callback, AsusLED.OFF, **kwargs)
 
     return True
diff --git a/tests/modules/test_led.py b/tests/modules/test_led.py
new file mode 100644
index 0000000..8ab4e81
--- /dev/null
+++ b/tests/modules/test_led.py
@@ -0,0 +1,67 @@
+"""Tests for the led module."""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from asusrouter.modules.endpoint import Endpoint
+from asusrouter.modules.identity import AsusDevice
+from asusrouter.modules.led import AsusLED, keep_state, set_state
+
+
+@pytest.mark.asyncio
+async def test_set_state():
+    """Test set_state."""
+
+    # Arrange
+    callback = AsyncMock(return_value=True)
+    state = AsusLED.ON
+    expect_modify = False
+
+    # Act
+    result = await set_state(callback, state, expect_modify=expect_modify)
+
+    # Assert
+    callback.assert_called_once_with(
+        service="start_ctrl_led",
+        arguments={"led_val": state.value},
+        apply=True,
+        expect_modify=expect_modify,
+    )
+    assert result is True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "identity, state, expected, set_state_calls",
+    [
+        (None, AsusLED.OFF, False, 0),
+        (MagicMock(spec=AsusDevice, endpoints={}), AsusLED.OFF, False, 0),
+        (
+            MagicMock(spec=AsusDevice, endpoints={Endpoint.SYSINFO: None}),
+            AsusLED.ON,
+            False,
+            0,
+        ),
+        (
+            MagicMock(spec=AsusDevice, endpoints={Endpoint.SYSINFO: "sysinfo"}),
+            AsusLED.OFF,
+            True,
+            2,
+        ),
+    ],
+)
+async def test_keep_state(identity, state, expected, set_state_calls):
+    """Test keep_state."""
+
+    # Arrange
+    callback = AsyncMock(return_value=True)
+    with patch(
+        "asusrouter.modules.led.set_state", new_callable=AsyncMock
+    ) as mock_set_state:
+        # Act
+        result = await keep_state(callback, state, identity=identity)
+
+        # Assert
+        assert result is expected
+        assert mock_set_state.call_count == set_state_calls