feat: add weight update support & fix behavior sensor device conflict…

custom_components/tryfi/number.py

@@ -14,6 +14,7 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfMass
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
@@ -106,10 +107,17 @@ def device_info(self) -> dict[str, Any]:
         return device_info
 
     async def async_set_native_value(self, value: float) -> None:
-        """Update the weight value."""
-        # Note: The API doesn't appear to have a method to update weight
-        # This would need to be implemented if the API supports it
-        _LOGGER.warning(
-            "Weight update not implemented - API method needed for pet %s",
-            self.pet.name if self.pet else "unknown",
-        )
+        """Update the weight value via the TryFi API."""
+        pet = self.pet
+        if not pet:
+            _LOGGER.warning("Cannot set weight: pet not available")
+            return
+
+        try:
+            await self.hass.async_add_executor_job(
+                pet.setWeight, self.coordinator.data.session, float(value)
+            )
+            await self.coordinator.async_request_refresh()
+        except Exception as err:
+            _LOGGER.error("Failed to set weight for %s: %s", pet.name, err)
+            raise HomeAssistantError(f"Failed to set weight for {pet.name}") from err

custom_components/tryfi/pytryfi/common/query.py

@@ -68,6 +68,7 @@
 
 FRAGMENT_WIFI_NETWORK_DETAILS = "fragment WifiNetworkDetails on WifiNetwork {  ssid  state  addressLabel  isHidden  position {    latitude    longitude  }}"
 QUERY_GET_WIFI_NETWORKS = "query GetWifiNetworks($householdId: ID!) {  household(id: $householdId) {    id    wifiNetworks {      credentialPackHash      maximumNetworkCount      networks {        __typename        ...WifiNetworkDetails      }    }  }}"
+MUTATION_UPDATE_PET = "mutation UpdatePet($input: UpdatePetInput!) { updatePet(input: $input) { __typename ...BasePetProfile } }"
 MUTATION_UPDATE_WIFI_NETWORK = "mutation UpdateWifiNetwork($input: UpdateWifiNetworkInput!) {  updateWifiNetwork(input: $input) {    __typename    ...WifiNetworkDetails  }}"
 
 REQUEST_FRAGMENTS_PET_ALL_INFO = (
@@ -200,6 +201,18 @@ def getPetHealthTrends(session: requests.Session, petId: str, period: str = "DAY
     return response["data"]["getPetHealthTrendsForPet"]
 
 
+def updatePetWeight(session: requests.Session, petId: str, weight_kg: float):
+    qString = MUTATION_UPDATE_PET + FRAGMENT_BASE_PET_PROFILE + FRAGMENT_BREED_DETAILS + FRAGMENT_PHOTO_DETAILS
+    qVariables = {
+        "input": {
+            "id": petId,
+            "weight": weight_kg
+        }
+    }
+    response = mutation(session, qString, qVariables)
+    LOGGER.debug(f"updatePetWeight: {response}")
+    return response['data']['updatePet']['weight']
+
 def setLedColor(session: requests.Session, deviceId: str, ledColorCode):
     qString = (
         MUTATION_SET_LED_COLOR

custom_components/tryfi/pytryfi/fiPet.py

@@ -325,6 +325,15 @@ def setBehaviorStatsFromTrends(self, behaviorTrends, period: str = "DAY"):
                 setattr(self, f"_{prefix}{attr_suffix}Count", events_count)
                 setattr(self, f"_{prefix}{attr_suffix}Duration", duration_minutes)
 
+    def setWeight(self, session: requests.Session, weight_kg: float):
+        try:
+            updated_weight = query.updatePetWeight(session, self.petId, weight_kg)
+            self._weight = float(updated_weight)
+            return True
+        except Exception as e:
+            LOGGER.error(f"Could not update weight for Pet {self.name}.\n{e}")
+            return False
+
     # set the color code of the led light on the pet collar
     def setLedColorCode(self, session: requests.Session, colorCode):
         try:

custom_components/tryfi/sensor.py

@@ -700,12 +700,10 @@ def __init__(
             self._attr_device_class = SensorDeviceClass.DURATION
             self._attr_state_class = SensorStateClass.MEASUREMENT
 
-        # Set device info
+        # Set device info — only identifiers, so the device registry
+        # merges with the main pet device instead of overriding model/name
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, pet.petId)},
-            name=pet.name,
-            manufacturer="TryFi",
-            model="Series 3+ Collar",
         )
 
     def _fipet_attr_name(self) -> str:

tests/pytryfi/test_pet.py

@@ -1,3 +1,5 @@
+from unittest.mock import patch
+
 from custom_components.tryfi.pytryfi import FiPet, FiDevice
 from .utils import mock_graphql, GRAPHQL_FIXTURE_PET_ALL_INFO, REQ_PET_ALL_INFO
 
@@ -118,3 +120,35 @@ def test_update_behavior_stats():
     assert pet.dailyLickingDuration == 6
     assert pet.dailyScratchingCount == 4
     assert pet.dailyScratchingDuration == 1
+
+
+@responses.activate
+def test_set_weight_success():
+    """Test setWeight updates the pet's weight via the API."""
+    responses.add(
+        method=responses.POST,
+        url="https://api.tryfi.com/graphql",
+        status=200,
+        json={"data": {"updatePet": {"__typename": "BasePet", "weight": 14.2}}},
+    )
+
+    pet = FiPet("test-pet")
+    pet._name = "Buddy"
+    result = pet.setWeight(requests.Session(), 14.2)
+
+    assert result is True
+    assert pet.weight == 14.2
+
+
+def test_set_weight_failure():
+    """Test setWeight returns False on API failure."""
+    pet = FiPet("test-pet")
+    pet._name = "Buddy"
+
+    with patch(
+        "custom_components.tryfi.pytryfi.common.query.updatePetWeight",
+        side_effect=Exception("API error"),
+    ):
+        result = pet.setWeight(requests.Session(), 10.0)
+
+    assert result is False

tests/pytryfi/test_query.py

@@ -8,7 +8,7 @@
 import json
 
 from custom_components.tryfi.pytryfi.exceptions import RemoteApiError
-from custom_components.tryfi.pytryfi.common.query import query
+from custom_components.tryfi.pytryfi.common.query import query, updatePetWeight
 from tests.pytryfi.utils import mock_graphql, mock_response
 
 
@@ -67,3 +67,22 @@ def test_query_graphql_errors():
 
     assert "GraphQL error" in str(exc_info.value)
     assert "Invalid query" in str(exc_info.value)
+
+
+@responses.activate
+def test_update_pet_weight():
+    """Test updatePetWeight sends mutation and returns updated weight."""
+    responses.add(
+        method=responses.POST,
+        url="https://api.tryfi.com/graphql",
+        status=200,
+        json={"data": {"updatePet": {"__typename": "BasePet", "weight": 15.5}}},
+    )
+
+    result = updatePetWeight(requests.Session(), "pet-123", 15.5)
+    assert result == 15.5
+
+    body = json.loads(responses.calls[0].request.body)
+    assert body["variables"]["input"]["id"] == "pet-123"
+    assert body["variables"]["input"]["weight"] == 15.5
+    assert "UpdatePetInput" in body["query"]

tests/test_number.py

@@ -1,8 +1,8 @@
-"""Test TryFi sensor platform."""
+"""Test TryFi number platform."""
 
 from __future__ import annotations
 
-from unittest.mock import Mock
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 
@@ -16,6 +16,7 @@ def mock_coordinator():
     """Create a mock coordinator."""
     coordinator = Mock()
     coordinator.data = Mock()
+    coordinator.async_request_refresh = AsyncMock()
     return coordinator
 
 
@@ -45,9 +46,65 @@ async def test_weight_number(
     mock_coordinator,
     mock_pet_with_stats,
 ) -> None:
-    """Test TryFi battery sensor."""
+    """Test TryFi weight number entity."""
     mock_coordinator.data.getPet.return_value = mock_pet_with_stats
 
     sensor = TryFiPetWeightNumber(mock_coordinator, mock_pet_with_stats)
 
     assert sensor.native_unit_of_measurement == "kg"
+
+
+async def test_async_set_native_value(
+    hass: HomeAssistant,
+    mock_coordinator,
+    mock_pet_with_stats,
+) -> None:
+    """Test async_set_native_value calls the API and refreshes coordinator."""
+    mock_coordinator.data.getPet.return_value = mock_pet_with_stats
+    mock_pet_with_stats.setWeight.return_value = True
+
+    entity = TryFiPetWeightNumber(mock_coordinator, mock_pet_with_stats)
+    entity.hass = hass
+
+    await entity.async_set_native_value(15.5)
+
+    mock_pet_with_stats.setWeight.assert_called_once_with(
+        mock_coordinator.data.session, 15.5
+    )
+    mock_coordinator.async_request_refresh.assert_awaited_once()
+
+
+async def test_async_set_native_value_no_pet(
+    hass: HomeAssistant,
+    mock_coordinator,
+    mock_pet_with_stats,
+) -> None:
+    """Test async_set_native_value when pet is not available."""
+    mock_coordinator.data.getPet.return_value = None
+
+    entity = TryFiPetWeightNumber(mock_coordinator, mock_pet_with_stats)
+    entity.hass = hass
+
+    await entity.async_set_native_value(15.5)
+
+    mock_coordinator.async_request_refresh.assert_not_awaited()
+
+
+async def test_async_set_native_value_api_error(
+    hass: HomeAssistant,
+    mock_coordinator,
+    mock_pet_with_stats,
+) -> None:
+    """Test async_set_native_value raises HomeAssistantError when API call fails."""
+    from homeassistant.exceptions import HomeAssistantError
+
+    mock_coordinator.data.getPet.return_value = mock_pet_with_stats
+    mock_pet_with_stats.setWeight.side_effect = ConnectionError("API error")
+
+    entity = TryFiPetWeightNumber(mock_coordinator, mock_pet_with_stats)
+    entity.hass = hass
+
+    with pytest.raises(HomeAssistantError, match="Failed to set weight for Fido"):
+        await entity.async_set_native_value(15.5)
+
+    mock_coordinator.async_request_refresh.assert_not_awaited()