Improve seventeentrack (#25454)

* Code improvement + tests

* review

* review + moving to pytest test function

* move test to async

* remove code comment

homeassistant/components/seventeentrack/sensor.py

@@ -33,11 +33,14 @@ DATA_SUMMARY = 'summary_data'
 DEFAULT_ATTRIBUTION = 'Data provided by 17track.net'
 DEFAULT_SCAN_INTERVAL = timedelta(minutes=10)
 
-ENTITY_ID_TEMPLATE = 'package_{0}_{1}'
+UNIQUE_ID_TEMPLATE = 'package_{0}_{1}'
+ENTITY_ID_TEMPLATE = 'sensor.seventeentrack_package_{0}'
 
-NOTIFICATION_DELIVERED_ID_SCAFFOLD = 'package_delivered_{0}'
+NOTIFICATION_DELIVERED_ID = 'package_delivered_{0}'
-NOTIFICATION_DELIVERED_TITLE = 'Package Delivered'
+NOTIFICATION_DELIVERED_TITLE = 'Package {0} delivered'
-NOTIFICATION_DELIVERED_URL_SCAFFOLD = 'https://t.17track.net/track#nums={0}'
+NOTIFICATION_DELIVERED_MESSAGE = 'Package Delivered: {0}<br />' + \
+                                 'Visit 17.track for more information: ' \
+                                 'https://t.17track.net/track#nums={1}'
 
 VALUE_DELIVERED = 'Delivered'
 
@@ -72,21 +75,11 @@ async def async_setup_platform(
 
     scan_interval = config.get(CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL)
 
-    data = SeventeenTrackData(
+    data = SeventeenTrackData(client, async_add_entities, scan_interval,
-        hass, client, async_add_entities, scan_interval,
+                              config[CONF_SHOW_ARCHIVED],
-        config[CONF_SHOW_ARCHIVED], config[CONF_SHOW_DELIVERED])
+                              config[CONF_SHOW_DELIVERED])
     await data.async_update()
 
-    sensors = []
-
-    for status, quantity in data.summary.items():
-        sensors.append(SeventeenTrackSummarySensor(data, status, quantity))
-
-    for package in data.packages:
-        sensors.append(SeventeenTrackPackageSensor(data, package))
-
-    async_add_entities(sensors, True)
-
 
 class SeventeenTrackSummarySensor(Entity):
     """Define a summary sensor."""
@@ -139,7 +132,7 @@ class SeventeenTrackSummarySensor(Entity):
         await self._data.async_update()
 
         package_data = []
-        for package in self._data.packages:
+        for package in self._data.packages.values():
             if package.status != self._status:
                 continue
 
@@ -175,14 +168,12 @@ class SeventeenTrackPackageSensor(Entity):
         self._friendly_name = package.friendly_name
         self._state = package.status
         self._tracking_number = package.tracking_number
+        self.entity_id = ENTITY_ID_TEMPLATE.format(self._tracking_number)
 
     @property
     def available(self):
         """Return whether the entity is available."""
-        return bool([
+        return self._data.packages.get(self._tracking_number) is not None
-            p for p in self._data.packages
-            if p.tracking_number == self._tracking_number
-        ])
 
     @property
     def device_state_attributes(self):
@@ -210,44 +201,24 @@ class SeventeenTrackPackageSensor(Entity):
     @property
     def unique_id(self):
         """Return a unique, HASS-friendly identifier for this entity."""
-        return ENTITY_ID_TEMPLATE.format(
+        return UNIQUE_ID_TEMPLATE.format(
             self._data.account_id, self._tracking_number)
 
     async def async_update(self):
         """Update the sensor."""
         await self._data.async_update()
 
-        if not self._data.packages:
+        if not self.available:
+            self.hass.async_create_task(self._remove())
             return
 
-        try:
+        package = self._data.packages.get(self._tracking_number, None)
-            package = next((
-                p for p in self._data.packages
-                if p.tracking_number == self._tracking_number))
-        except StopIteration:
-            # If the package no longer exists in the data, log a message and
-            # delete this entity:
-            _LOGGER.info(
-                'Deleting entity for stale package: %s', self._tracking_number)
-            reg = await self.hass.helpers.entity_registry.async_get_registry()
-            self.hass.async_create_task(reg.async_remove(self.entity_id))
-            self.hass.async_create_task(self.async_remove())
-            return
 
         # If the user has elected to not see delivered packages and one gets
         # delivered, post a notification:
         if package.status == VALUE_DELIVERED and not self._data.show_delivered:
-            _LOGGER.info('Package delivered: %s', self._tracking_number)
+            self._notify_delivered()
-            self.hass.components.persistent_notification.create(
+            self.hass.async_create_task(self._remove())
-                'Package Delivered: {0}<br />'
-                'Visit 17.track for more infomation: {1}'
-                ''.format(
-                    self._tracking_number,
-                    NOTIFICATION_DELIVERED_URL_SCAFFOLD.format(
-                        self._tracking_number)),
-                title=NOTIFICATION_DELIVERED_TITLE,
-                notification_id=NOTIFICATION_DELIVERED_ID_SCAFFOLD.format(
-                    self._tracking_number))
             return
 
         self._attrs.update({
@@ -255,26 +226,53 @@ class SeventeenTrackPackageSensor(Entity):
             ATTR_LOCATION: package.location,
         })
         self._state = package.status
+        self._friendly_name = package.friendly_name
+
+    async def _remove(self):
+        """Remove entity itself."""
+        await self.async_remove()
+
+        reg = await self.hass.helpers.entity_registry.async_get_registry()
+        entity_id = reg.async_get_entity_id(
+            'sensor', 'seventeentrack',
+            UNIQUE_ID_TEMPLATE.format(
+                self._data.account_id, self._tracking_number))
+        if entity_id:
+            reg.async_remove(entity_id)
+
+    def _notify_delivered(self):
+        """Notify when package is delivered."""
+        _LOGGER.info('Package delivered: %s', self._tracking_number)
+
+        identification = self._friendly_name if self._friendly_name else \
+            self._tracking_number
+        message = NOTIFICATION_DELIVERED_MESSAGE.format(self._tracking_number,
+                                                        identification)
+        title = NOTIFICATION_DELIVERED_TITLE.format(identification)
+        notification_id = NOTIFICATION_DELIVERED_TITLE\
+            .format(self._tracking_number)
+
+        self.hass.components.persistent_notification\
+            .create(message, title=title, notification_id=notification_id)
 
 
 class SeventeenTrackData:
     """Define a data handler for 17track.net."""
 
-    def __init__(
+    def __init__(self, client, async_add_entities, scan_interval,
-            self, hass, client, async_add_entities, scan_interval,
+                 show_archived, show_delivered):
-            show_archived, show_delivered):
         """Initialize."""
         self._async_add_entities = async_add_entities
         self._client = client
-        self._hass = hass
         self._scan_interval = scan_interval
         self._show_archived = show_archived
         self.account_id = client.profile.account_id
-        self.packages = []
+        self.packages = {}
         self.show_delivered = show_delivered
         self.summary = {}
 
         self.async_update = Throttle(self._scan_interval)(self._async_update)
+        self.first_update = True
 
     async def _async_update(self):
         """Get updated data from 17track.net."""
@@ -285,46 +283,36 @@ class SeventeenTrackData:
                 show_archived=self._show_archived)
             _LOGGER.debug('New package data received: %s', packages)
 
-            if not self.show_delivered:
+            new_packages = {p.tracking_number: p for p in packages}
-                packages = [p for p in packages if p.status != VALUE_DELIVERED]
 
-            packages_map = {p.tracking_number: p for p in packages}
+            to_add = set(new_packages) - set(self.packages)
-            # Add new packages:
 
-            already_tracked_nbr = [p.tracking_number for p in self.packages]
-            received_tracked_nbr = [p.tracking_number for p in packages]
-
-            to_add = set(received_tracked_nbr) - set(already_tracked_nbr)
             _LOGGER.debug('Will add new tracking numbers: %s', to_add)
-            if self.packages and to_add:
+            if to_add:
                 self._async_add_entities([
                     SeventeenTrackPackageSensor(self,
-                                                packages_map[tracking_number])
+                                                new_packages[tracking_number])
                     for tracking_number in to_add
                 ], True)
 
-            # Remove archived packages from the entity registry:
+            self.packages = new_packages
-            to_remove = set(received_tracked_nbr) - set(already_tracked_nbr)
-            _LOGGER.debug('Will remove tracking number: %s', to_remove)
-            reg = await self._hass.helpers.entity_registry.async_get_registry()
-            for tracking_number in to_remove:
-                entity_id = reg.async_get_entity_id(
-                    'sensor', 'seventeentrack',
-                    ENTITY_ID_TEMPLATE.format(
-                        self.account_id, tracking_number))
-                if not entity_id:
-                    continue
-                self._hass.async_create_task(reg.async_remove(entity_id))
-
-            self.packages = packages
         except SeventeenTrackError as err:
             _LOGGER.error('There was an error retrieving packages: %s', err)
-            self.packages = []
 
         try:
             self.summary = await self._client.profile.summary(
                 show_archived=self._show_archived)
             _LOGGER.debug('New summary data received: %s', self.summary)
+
+            # creating summary sensors on first update
+            if self.first_update:
+                self.first_update = False
+
+                self._async_add_entities([
+                    SeventeenTrackSummarySensor(self, status, quantity)
+                    for status, quantity in self.summary.items()
+                ], True)
+
         except SeventeenTrackError as err:
             _LOGGER.error('There was an error retrieving the summary: %s', err)
             self.summary = {}

requirements_test_all.txt

@@ -240,6 +240,9 @@ pushbullet.py==0.11.0
 # homeassistant.components.canary
 py-canary==0.5.0
 
+# homeassistant.components.seventeentrack
+py17track==2.2.2
+
 # homeassistant.components.tplink
 pyHS100==0.3.5
 

script/gen_requirements_all.py

@@ -161,6 +161,7 @@ TEST_REQUIREMENTS = (
     'zeroconf',
     'zigpy-homeassistant',
     'bellows-homeassistant',
+    'py17track',
 )
 
 IGNORE_PIN = ('colorlog>2.1,<3', 'keyring>=9.3,<10.0', 'urllib3')

tests/components/seventeentrack/__init__.py

@@ -0,0 +1 @@
+"""Tests for the seventeentrack component."""

tests/components/seventeentrack/test_sensor.py

@@ -0,0 +1,311 @@
+"""Tests for the seventeentrack sensor."""
+import datetime
+from typing import Union
+
+import pytest
+import mock
+from py17track.package import Package
+
+from homeassistant.components.seventeentrack.sensor \
+    import CONF_SHOW_ARCHIVED, CONF_SHOW_DELIVERED
+from homeassistant.const import CONF_USERNAME, CONF_PASSWORD
+from homeassistant.setup import async_setup_component
+from homeassistant.util import utcnow
+from tests.common import MockDependency, async_fire_time_changed
+
+VALID_CONFIG_MINIMAL = {
+    'sensor': {
+        'platform': 'seventeentrack',
+        CONF_USERNAME: 'test',
+        CONF_PASSWORD: 'test'
+    }
+}
+
+INVALID_CONFIG = {
+    'sensor': {
+        'platform': 'seventeentrack',
+        'boom': 'test',
+    }
+}
+
+VALID_CONFIG_FULL = {
+    'sensor': {
+        'platform': 'seventeentrack',
+        CONF_USERNAME: 'test',
+        CONF_PASSWORD: 'test',
+        CONF_SHOW_ARCHIVED: True,
+        CONF_SHOW_DELIVERED: True
+    }
+}
+
+VALID_CONFIG_FULL_NO_DELIVERED = {
+    'sensor': {
+        'platform': 'seventeentrack',
+        CONF_USERNAME: 'test',
+        CONF_PASSWORD: 'test',
+        CONF_SHOW_ARCHIVED: False,
+        CONF_SHOW_DELIVERED: False
+    }
+}
+
+DEFAULT_SUMMARY = {
+    "Not Found": 0,
+    "In Transit": 0,
+    "Expired": 0,
+    "Ready to be Picked Up": 0,
+    "Undelivered": 0,
+    "Delivered": 0,
+    "Returned": 0
+}
+
+NEW_SUMMARY_DATA = {
+    "Not Found": 1,
+    "In Transit": 1,
+    "Expired": 1,
+    "Ready to be Picked Up": 1,
+    "Undelivered": 1,
+    "Delivered": 1,
+    "Returned": 1
+}
+
+
+class ClientMock:
+    """Mock the py17track client to inject the ProfileMock."""
+
+    def __init__(self, websession) -> None:
+        """Mock the profile."""
+        self.profile = ProfileMock()
+
+
+class ProfileMock:
+    """ProfileMock will mock data coming from 17track."""
+
+    package_list = []
+    login_result = True
+    summary_data = DEFAULT_SUMMARY
+    account_id = '123'
+
+    @classmethod
+    def reset(cls):
+        """Reset data to defaults."""
+        cls.package_list = []
+        cls.login_result = True
+        cls.summary_data = DEFAULT_SUMMARY
+        cls.account_id = '123'
+
+    def __init__(self) -> None:
+        """Override Account id."""
+        self.account_id = self.__class__.account_id
+
+    async def login(self, email: str, password: str) -> bool:
+        """Login mock."""
+        return self.__class__.login_result
+
+    async def packages(self, package_state: Union[int, str] = '',
+                       show_archived: bool = False) -> list:
+        """Packages mock."""
+        return self.__class__.package_list[:]
+
+    async def summary(self, show_archived: bool = False) -> dict:
+        """Summary mock."""
+        return self.__class__.summary_data
+
+
+@pytest.fixture(autouse=True, name="mock_py17track")
+def fixture_mock_py17track():
+    """Mock py17track dependency."""
+    with MockDependency('py17track'):
+        yield
+
+
+@pytest.fixture(autouse=True, name="mock_client")
+def fixture_mock_client(mock_py17track):
+    """Mock py17track client."""
+    with mock.patch('py17track.Client', new=ClientMock):
+        yield
+    ProfileMock.reset()
+
+
+async def _setup_seventeentrack(hass, config=None, summary_data=None):
+    """Set up component using config."""
+    if not config:
+        config = VALID_CONFIG_MINIMAL
+    if not summary_data:
+        summary_data = {}
+
+    ProfileMock.summary_data = summary_data
+    assert await async_setup_component(hass, 'sensor', config)
+
+
+async def _goto_future(hass, future=None):
+    """Move to future."""
+    if not future:
+        future = utcnow() + datetime.timedelta(minutes=10)
+    with mock.patch('homeassistant.util.utcnow', return_value=future):
+        async_fire_time_changed(hass, future)
+        await hass.async_block_till_done()
+
+
+async def test_full_valid_config(hass):
+    """Ensure everything starts correctly."""
+    assert await async_setup_component(hass, 'sensor', VALID_CONFIG_FULL)
+    assert len(hass.states.async_entity_ids()) == len(
+        ProfileMock.summary_data.keys())
+
+
+async def test_valid_config(hass):
+    """Ensure everything starts correctly."""
+    assert await async_setup_component(hass, 'sensor', VALID_CONFIG_MINIMAL)
+
+    assert len(hass.states.async_entity_ids()) == len(
+        ProfileMock.summary_data.keys())
+
+
+async def test_invalid_config(hass):
+    """Ensure nothing is created when config is wrong."""
+    assert await async_setup_component(hass, 'sensor', INVALID_CONFIG)
+
+    assert not hass.states.async_entity_ids()
+
+
+async def test_add_package(hass):
+    """Ensure package is added correctly when user add a new package."""
+    package = Package('456', 206, 'friendly name 1', 'info text 1',
+                      'location 1', 206, 2)
+    ProfileMock.package_list = [package]
+
+    await _setup_seventeentrack(hass)
+    assert hass.states.get(
+        'sensor.seventeentrack_package_456') is not None
+    assert len(hass.states.async_entity_ids()) == 1
+
+    package2 = Package('789', 206, 'friendly name 2', 'info text 2',
+                       'location 2', 206, 2)
+    ProfileMock.package_list = [package, package2]
+
+    await _goto_future(hass)
+
+    assert hass.states.get(
+        'sensor.seventeentrack_package_789') is not None
+    assert len(hass.states.async_entity_ids()) == 2
+
+
+async def test_remove_package(hass):
+    """Ensure entity is not there anymore if package is not there."""
+    package1 = Package('456', 206, 'friendly name 1', 'info text 1',
+                       'location 1', 206, 2)
+    package2 = Package('789', 206, 'friendly name 2', 'info text 2',
+                       'location 2', 206, 2)
+
+    ProfileMock.package_list = [package1, package2]
+
+    await _setup_seventeentrack(hass)
+
+    assert hass.states.get(
+        'sensor.seventeentrack_package_456') is not None
+    assert hass.states.get(
+        'sensor.seventeentrack_package_789') is not None
+    assert len(hass.states.async_entity_ids()) == 2
+
+    ProfileMock.package_list = [package2]
+
+    await _goto_future(hass)
+
+    assert hass.states.get(
+        'sensor.seventeentrack_package_456') is None
+    assert hass.states.get(
+        'sensor.seventeentrack_package_789') is not None
+    assert len(hass.states.async_entity_ids()) == 1
+
+
+async def test_friendly_name_changed(hass):
+    """Test friendly name change."""
+    package = Package('456', 206, 'friendly name 1', 'info text 1',
+                      'location 1', 206, 2)
+    ProfileMock.package_list = [package]
+
+    await _setup_seventeentrack(hass)
+
+    assert hass.states.get(
+        'sensor.seventeentrack_package_456') is not None
+    assert len(hass.states.async_entity_ids()) == 1
+
+    package = Package('456', 206, 'friendly name 2', 'info text 1',
+                      'location 1', 206, 2)
+    ProfileMock.package_list = [package]
+
+    await _goto_future(hass)
+
+    assert hass.states.get(
+        'sensor.seventeentrack_package_456') is not None
+    entity = hass.data['entity_components']['sensor'].get_entity(
+        'sensor.seventeentrack_package_456')
+    assert entity.name == 'Seventeentrack Package: friendly name 2'
+    assert len(hass.states.async_entity_ids()) == 1
+
+
+async def test_delivered_not_shown(hass):
+    """Ensure delivered packages are not shown."""
+    package = Package('456', 206, 'friendly name 1', 'info text 1',
+                      'location 1', 206, 2, 40)
+    ProfileMock.package_list = [package]
+
+    hass.components.persistent_notification = mock.MagicMock()
+    await _setup_seventeentrack(hass, VALID_CONFIG_FULL_NO_DELIVERED)
+    assert not hass.states.async_entity_ids()
+    hass.components.persistent_notification.create.assert_called()
+
+
+async def test_delivered_shown(hass):
+    """Ensure delivered packages are show when user choose to show them."""
+    package = Package('456', 206, 'friendly name 1', 'info text 1',
+                      'location 1', 206, 2, 40)
+    ProfileMock.package_list = [package]
+
+    hass.components.persistent_notification = mock.MagicMock()
+    await _setup_seventeentrack(hass, VALID_CONFIG_FULL)
+
+    assert hass.states.get(
+        'sensor.seventeentrack_package_456') is not None
+    assert len(hass.states.async_entity_ids()) == 1
+    hass.components.persistent_notification.create.assert_not_called()
+
+
+async def test_becomes_delivered_not_shown_notification(hass):
+    """Ensure notification is triggered when package becomes delivered."""
+    package = Package('456', 206, 'friendly name 1', 'info text 1',
+                      'location 1', 206, 2)
+    ProfileMock.package_list = [package]
+
+    await _setup_seventeentrack(hass, VALID_CONFIG_FULL_NO_DELIVERED)
+
+    assert hass.states.get(
+        'sensor.seventeentrack_package_456') is not None
+    assert len(hass.states.async_entity_ids()) == 1
+
+    package_delivered = Package('456', 206, 'friendly name 1',
+                                'info text 1', 'location 1', 206, 2, 40)
+    ProfileMock.package_list = [package_delivered]
+
+    hass.components.persistent_notification = mock.MagicMock()
+    await _goto_future(hass)
+
+    hass.components.persistent_notification.create.assert_called()
+    assert not hass.states.async_entity_ids()
+
+
+async def test_summary_correctly_updated(hass):
+    """Ensure summary entities are not duplicated."""
+    await _setup_seventeentrack(hass, summary_data=DEFAULT_SUMMARY)
+
+    assert len(hass.states.async_entity_ids()) == 7
+    for state in hass.states.async_all():
+        assert state.state == '0'
+
+    ProfileMock.summary_data = NEW_SUMMARY_DATA
+
+    await _goto_future(hass)
+
+    assert len(hass.states.async_entity_ids()) == 7
+    for state in hass.states.async_all():
+        assert state.state == '1'