Merge pull request #9 from id8-engineering/add_functions

Dexter9532 · web-flow · commit 22c5fa1814e0 · 2025-11-13T12:01:17.000+01:00
feat: add more functions
diff --git a/src/dcm230/dcm230.py b/src/dcm230/dcm230.py
@@ -12,7 +12,7 @@
 import struct
 from dataclasses import dataclass
 from decimal import Decimal
-from typing import Final, TypeVar
+from typing import ClassVar, Final, TypeVar
 
 from pymodbus.client import ModbusSerialClient
 from pymodbus.exceptions import ModbusException
@@ -54,7 +54,7 @@ def register_properties(cls: T) -> T:
     a corresponding @property getter, and optionally a setter if `writable=True`.
 
     The generated getter automatically calls `_read_register(register_name)`
-    and the setter calls `_write_register(address, value)` with range validation
+    and the setter calls `_write_registers(address, value)` with range validation
     if enabled in the `RegisterSpec`.
 
     Args:
@@ -97,7 +97,7 @@ def setter(self: "Dcm230", value: int, _name: str = name, _spec: "RegisterSpec"
             if _spec.range and not (_spec.min <= value <= _spec.max):
                 msg = f"Invalid value for '{_name}': {value}. Must be between {_spec.min} and {_spec.max}."
                 raise ValueError(msg)
-            self._write_register(_spec.address, int(value))
+            self._write_registers(_spec.address, int(value))
 
         prop = property(getter, setter) if spec.writable else property(getter)
 
@@ -121,11 +121,56 @@ class Dcm230:
 
     SINGLE_REGISTER = 1
     MAX_REGS = 2
+
     INPUT_REGISTER = 0x03
     HOLDING_REGISTER = 0x04
 
+    BACKLIT_OPTIONS: ClassVar[list[int]] = [0, 5, 10, 20, 30, 60]
+    DCM230_REGISTER_BACKLIT_TIME = 0x003C
+    DCM230_REGISTER_RESET_MAX_DMD_AND_PARTIAL_ENERGY = 0xF010
+    RESET_MAX_DMD = 0x0000
+    RESET_PARTIAL_ENERGY = 0x0003
+
     _register_specs: Final[dict[str, RegisterSpec]] = {
-        "V": RegisterSpec(address=0x0000, count=2, decimals=1, reg_type=0x03),
+        "V": RegisterSpec(address=0x0000, count=2, decimals=1, reg_type=INPUT_REGISTER),
+        "A": RegisterSpec(address=0x0006, count=2, decimals=1, reg_type=INPUT_REGISTER),
+        "W": RegisterSpec(address=0x000C, count=2, decimals=1, reg_type=INPUT_REGISTER),
+        "kwh": RegisterSpec(address=0x0048, count=2, decimals=1, reg_type=INPUT_REGISTER),
+        "W_dmd": RegisterSpec(address=0x0054, count=2, decimals=1, reg_type=INPUT_REGISTER),
+        "W_dmd_peak": RegisterSpec(address=0x0056, count=2, decimals=1, reg_type=INPUT_REGISTER),
+        "kwh_tot": RegisterSpec(address=0x0156, count=2, decimals=1, reg_type=INPUT_REGISTER),
+        "kwh_partial": RegisterSpec(address=0x0180, count=2, decimals=1, reg_type=INPUT_REGISTER),
+        "dmd_period": RegisterSpec(
+            address=0x0002,
+            count=2,
+            range=True,
+            min=0,
+            max=60,
+            writable=True,
+            reg_type=HOLDING_REGISTER,
+            return_type=int,
+        ),
+        "network_info": RegisterSpec(
+            address=0x0012, count=2, range=True, min=0, max=3, writable=True, reg_type=HOLDING_REGISTER, return_type=int
+        ),
+        "device_id": RegisterSpec(
+            address=0x0014,
+            count=2,
+            range=True,
+            min=1,
+            max=247,
+            writable=True,
+            reg_type=HOLDING_REGISTER,
+            return_type=int,
+        ),
+        "password": RegisterSpec(address=0x0018, count=2, reg_type=HOLDING_REGISTER, return_type=int),
+        "baud_rate": RegisterSpec(
+            address=0x001C, count=2, range=True, min=0, max=5, writable=True, reg_type=HOLDING_REGISTER, return_type=int
+        ),
+        "energy_measurement_tool": RegisterSpec(
+            address=0xF920, count=2, range=True, min=0, max=3, writable=True, reg_type=HOLDING_REGISTER, return_type=int
+        ),
+        "serial_number": RegisterSpec(address=0xFC00, count=2, reg_type=HOLDING_REGISTER, return_type=int),
     }
 
     def __init__(self, device_address: int, client: ModbusSerialClient) -> None:
@@ -190,8 +235,8 @@ def _read_register(self, register_name: str) -> Decimal | int:
             return round(value, spec.decimals)
         return self._unpack(regs, spec.address)
 
-    def _write_register(self, address: int, value: int) -> None:
-        """Write a single Modbus register.
+    def _write_registers(self, address: int, value: int) -> None:
+        """Write to Modbus registers.
 
         Args:
             address: Register address to write.
@@ -200,12 +245,9 @@ def _write_register(self, address: int, value: int) -> None:
         Raises:
             ModbusException: If the write operation fails.
         """
-        result = self.client.write_register(address=address, value=value, device_id=self.device_address)
+        result = self.client.write_registers(address=address, values=[value], device_id=self.device_address)
         if result.isError():
-            msg = (
-                "Failed to write to single register. "
-                f"device_address={self.device_address} address={address} value={value}"
-            )
+            msg = f"Failed to write to registers. device_address={self.device_address} address={address} value={value}"
             raise ModbusException(msg)
 
     def _unpack(self, regs: list[int], address: int) -> int:
@@ -232,3 +274,60 @@ def _unpack(self, regs: list[int], address: int) -> int:
             raise ValueError(msg)
 
         return struct.unpack(">f", struct.pack(">HH", regs[0], regs[1]))[0]
+
+    def reset_max_dmd(self) -> None:
+        """Reset max demand.
+
+        Raises:
+            ModbusException: If failed to write to registers.
+        """
+        self._write_registers(self.DCM230_REGISTER_RESET_MAX_DMD_AND_PARTIAL_ENERGY, self.RESET_MAX_DMD)
+
+    def reset_partial_energy(self) -> None:
+        """Reset partial energy.
+
+        Raises:
+            ModbusException: If failed to write to registers.
+        """
+        self._write_registers(self.DCM230_REGISTER_RESET_MAX_DMD_AND_PARTIAL_ENERGY, self.RESET_PARTIAL_ENERGY)
+
+    @property
+    def backlit_time(self) -> int:
+        """Backlit time.
+
+        Options:
+            0, 5, 10, 20, 30, 60 minutes.
+
+        Returns:
+            int: Current backlit time value.
+
+        Raises:
+            ValueError: If not a specified value.
+            ModbusException: If failed to read holding registers.
+        """
+        regs = self._read_registers(self.DCM230_REGISTER_BACKLIT_TIME, self.MAX_REGS, self.HOLDING_REGISTER)
+        value = round(self._unpack(regs, self.DCM230_REGISTER_BACKLIT_TIME))
+
+        if value not in self.BACKLIT_OPTIONS:
+            msg = f"Invalid backlit option: {value}. Must be one of: {self.BACKLIT_OPTIONS}"
+            raise ValueError(msg)
+        return value
+
+    @backlit_time.setter
+    def backlit_time(self, value: int) -> None:
+        """Backlite time.
+
+        Options:
+            0, 5, 10, 20, 30, 60 minutes.
+
+        Args:
+            value (int): Set backlit time.
+
+        Raises:
+            ModbusException: If failed to write to registers.
+            ValueError: If not specified value.
+        """
+        if value not in self.BACKLIT_OPTIONS:
+            msg = f"Invalid backlit option: {value}. Must be one of: {self.BACKLIT_OPTIONS}"
+            raise ValueError(msg)
+        self._write_registers(self.DCM230_REGISTER_BACKLIT_TIME, value)
diff --git a/src/dcm230/dcm230.pyi b/src/dcm230/dcm230.pyi
@@ -2,12 +2,50 @@ from decimal import Decimal
 
 from pymodbus.client import ModbusSerialClient
 
+class RegisterSpec:
+    address: int
+    count: int
+    reg_type: int
+    decimals: int = 0
+    range: bool = False
+    min: int = 0
+    max: int = 0x7FFFFFFF
+    writable: bool = False
+    return_type: type[int | Decimal] = ...
+
+    def register_properties(self) -> None: ...
+
 class Dcm230:
+    RESET_MAX_DMD: int
+    RESET_PARTIAL_ENERGY: int
+    DCM230_REGISTER_RESET_MAX_DMD_AND_PARTIAL_ENERGY: int
+    DCM230_REGISTER_BACKLIT_TIME: int
+    INPUT_REGISTER: int
+    HOLDING_REGISTER: int
+    _register_specs: RegisterSpec
+
     def __init__(self, device_address: int, client: ModbusSerialClient) -> None: ...
     V: Decimal
+    A: Decimal
+    W: Decimal
+    kwh: Decimal
+    W_dmd: Decimal
+    W_dmd_peak: Decimal
+    kwh_tot: Decimal
+    kwh_partial: Decimal
+    dmd_period: Decimal
+    backlit_time: int
+    network_info: int
+    device_id: int
+    password: int
+    baud_rate: int
+    energy_measurement_tool: int
+    serial_number: int
 
     def _unpack(self, registers: list[int], address: int) -> int: ...
-    def _write_register(self, address: int, value: int) -> None: ...
+    def _write_registers(self, address: int, value: int) -> None: ...
     def _read_register(self, register_name: str) -> Decimal | int: ...
     def _read_input_registers(self, address: int, count: int) -> list[int]: ...
     def _read_registers(self, address: int, count: int, reg_type: int) -> list[int]: ...
+    def reset_max_dmd(self) -> None: ...
+    def reset_partial_energy(self) -> None: ...
diff --git a/src/dcm230/test_dcm230.py b/src/dcm230/test_dcm230.py
@@ -1,4 +1,4 @@
-# ruff: noqa: S101,PLR2004, N802, SLF001
+# ruff: noqa: S101, PLR2004, SLF001
 
 """Test file for driver."""
 
@@ -42,21 +42,128 @@ def test_read_register() -> None:
         _ = meter._read_registers(address, count, reg_type)
 
 
-def test_V() -> None:
-    """Test get v."""
+def test_range_validation() -> None:
+    """Test range validation."""
+    client = MagicMock()
+    meter = Dcm230(1, client)
+
+    """Test 1: Should not pass due to out of range."""
+    for name, spec in Dcm230._register_specs.items():  # type: ignore[attr-defined]
+        if not spec.range:
+            continue
+
+        client.write_registers.return_value.isError.return_value = False
+        with nullcontext():
+            _ = setattr(meter, name, spec.min)
+
+        invalid_value = spec.max + 1
+        with pytest.raises(ValueError, match="Invalid value for"):
+            setattr(meter, name, invalid_value)
+
+        invalid_value = spec.min - 1
+        with pytest.raises(ValueError, match="Invalid value for"):
+            setattr(meter, name, invalid_value)
+
+
+def read_input_registers() -> None:
+    """Test all input registers."""
     client = MagicMock()
     mock_result = MagicMock()
     mock_result.isError.return_value = False
     meter = Dcm230(1, client)
 
-    """Test 1: should pass"""
-    mock_result.registers = [0x4366, 0x0000]
+    mock_result.registers = [0x3F80, 0x0000]
     client.read_input_registers.return_value = mock_result
-    value = meter.V
-    assert value == 230
 
-    """Test 2: Should pass."""
-    mock_result.registers = [0x4624, 0x1000]
-    client.read_input_registers.return_value = mock_result
-    value = meter.V
-    assert value == 10500
+    """Test 1: Should pass."""
+    for name, spec in Dcm230._register_specs.items():  # type: ignore[attr-defined]
+        if spec.reg_type != meter.INPUT_REGISTER:
+            continue
+
+        value = getattr(meter, name)
+        assert value == 1
+
+
+def read_holding_registers() -> None:
+    """Test all holding registers."""
+    client = MagicMock()
+    mock_result = MagicMock()
+    mock_result.isError.return_value = False
+    meter = Dcm230(1, client)
+
+    mock_result.registers = [0x3F80, 0x0000]
+    client.read_holding_registers.return_value = mock_result
+
+    """Test 1: Should pass."""
+    for name, spec in Dcm230._register_specs.items():  # type: ignore[attr-defined]
+        if spec.reg_type != meter.HOLDING_REGISTER:
+            continue
+
+        value = getattr(meter, name)
+        assert value == 1
+
+
+def test_set_all_register() -> None:
+    """Test set all registers."""
+    client = MagicMock()
+    mock_result = MagicMock()
+    mock_result.isError.return_value = False
+    client.write_registers.return_value = mock_result
+
+    meter = Dcm230(1, client)
+
+    for name, spec in Dcm230._register_specs.items():  # type: ignore[attr-defined]
+        if not spec.writable:
+            continue
+
+        value = 1
+        setattr(meter, name, value)
+        client.write_registers.assert_called_once_with(address=spec.address, values=[1], device_id=1)
+        client.write_registers.reset_mock()
+
+
+def test_backlit_time() -> None:
+    """Test get and set backlit time."""
+    client = MagicMock()
+    mock_result = MagicMock()
+    mock_result.isError.return_value = False
+    meter = Dcm230(1, client)
+
+    """Test 1: should pass"""
+    mock_result.registers = [0x4270, 0x0000]
+    client.read_holding_registers.return_value = mock_result
+    value = meter.backlit_time
+    assert value == 60
+
+    """Test 2: Should raise exception due to invalid backlit option."""
+    with pytest.raises(ValueError, match="Invalid backlit option:"):
+        meter.backlit_time = 100
+
+    """Test 3: should pass when set."""
+    client.write_registers.return_value = mock_result
+    meter.backlit_time = 60
+    client.write_registers.assert_called_once_with(address=meter.DCM230_REGISTER_BACKLIT_TIME, values=[60], device_id=1)
+
+
+def test_reset_functions() -> None:
+    """Test all reset functions."""
+    client = MagicMock()
+    mock_result = MagicMock()
+    mock_result.isError.return_value = False
+    client.write_registers.return_value = mock_result
+
+    meter = Dcm230(1, client)
+
+    """Test 1: reset_max_dmd."""
+    meter.reset_max_dmd()
+    client.write_registers.assert_called_once_with(
+        address=meter.DCM230_REGISTER_RESET_MAX_DMD_AND_PARTIAL_ENERGY, values=[meter.RESET_MAX_DMD], device_id=1
+    )
+
+    client.write_registers.reset_mock()
+
+    """Test 2: reset_partial_energy."""
+    meter.reset_partial_energy()
+    client.write_registers.assert_called_once_with(
+        address=meter.DCM230_REGISTER_RESET_MAX_DMD_AND_PARTIAL_ENERGY, values=[meter.RESET_PARTIAL_ENERGY], device_id=1
+    )
