feat: support Canda Post mail count (#1068)

custom_components/mail_and_packages/const.py

@@ -38,6 +38,7 @@
 ATTR_EMAIL = "email"
 ATTR_SUBJECT = "subject"
 ATTR_BODY = "body"
+ATTR_BODY_COUNT = "body_count"
 ATTR_PATTERN = "pattern"
 ATTR_USPS_MAIL = "usps_mail"
 
@@ -290,6 +291,12 @@
     "capost_delivering": {},
     "capost_packages": {},
     "capost_tracking": {},
+    "capost_mail": {
+        "email": ["donotreply-nepasrepondre@communications.canadapost-postescanada.ca"],
+        "subject": ["You have mail on the way"],
+        "body": ["\\sYou have (\\d) pieces of mail\\s"],
+        "body_count": True,
+    },
     # DHL
     "dhl_delivered": {
         "email": [
@@ -838,6 +845,12 @@
         icon="mdi:truck-delivery",
         key="capost_delivering",
     ),
+    "capost_mail": SensorEntityDescription(
+        name="Mail Canada Post Mail",
+        native_unit_of_measurement="piece(s)",
+        icon="mdi:mailbox-up",
+        key="capost_mail",
+    ),
     "capost_packages": SensorEntityDescription(
         name="Mail Canada Post Packages",
         native_unit_of_measurement="package(s)",

custom_components/mail_and_packages/helpers.py

@@ -58,6 +58,7 @@
     AMAZON_TIME_PATTERN_END,
     ATTR_AMAZON_IMAGE,
     ATTR_BODY,
+    ATTR_BODY_COUNT,
     ATTR_CODE,
     ATTR_COUNT,
     ATTR_EMAIL,
@@ -1057,7 +1058,11 @@ def get_count(
         )
         if server_response == "OK" and data[0] is not None:
             if ATTR_BODY in SENSOR_DATA[sensor_type].keys():
-                count += find_text(data, account, SENSOR_DATA[sensor_type][ATTR_BODY])
+                body_count = SENSOR_DATA[sensor_type].get(ATTR_BODY_COUNT, False)
+                _LOGGER.debug("Check body for mail count? %s", body_count)
+                count += find_text(
+                    data, account, SENSOR_DATA[sensor_type][ATTR_BODY], body_count
+                )
             else:
                 count += len(data[0].split())
 
@@ -1151,7 +1156,9 @@ def get_tracking(
     return tracking
 
 
-def find_text(sdata: Any, account: Type[imaplib.IMAP4_SSL], search_terms: list) -> int:
+def find_text(
+    sdata: Any, account: Type[imaplib.IMAP4_SSL], search_terms: list, count: bool
+) -> int:
     """Filter for specific words in email.
 
     Return count of items found as integer
@@ -1175,7 +1182,18 @@ def find_text(sdata: Any, account: Type[imaplib.IMAP4_SSL], search_terms: list)
                         email_msg = part.get_payload(decode=True)
                         email_msg = email_msg.decode("utf-8", "ignore")
                         pattern = re.compile(rf"{search}")
-                        if (found := pattern.findall(email_msg)) and len(found) > 0:
+                        if (
+                            count
+                            and (found := pattern.search(email_msg))
+                            and len(found.groups()) > 0
+                        ):
+                            _LOGGER.debug(
+                                "Found (%s) in email result: %s",
+                                search,
+                                str(found.groups()),
+                            )
+                            count = int(found.group(1))
+                        elif (found := pattern.findall(email_msg)) and len(found) > 0:
                             _LOGGER.debug(
                                 "Found (%s) in email %s times.", search, str(len(found))
                             )

tests/conftest.py

@@ -22,6 +22,7 @@
 from tests.const import (
     FAKE_CONFIG_DATA,
     FAKE_CONFIG_DATA_AMAZON_FWD_STRING,
+    FAKE_CONFIG_DATA_CAPOST,
     FAKE_CONFIG_DATA_CUSTOM_IMG,
     FAKE_CONFIG_DATA_EXTERNAL,
     FAKE_CONFIG_DATA_MISSING_TIMEOUT,
@@ -164,6 +165,22 @@ async def integration_fixture_7(hass):
     return entry
 
 
+@pytest.fixture(name="integration_capost")
+async def integration_fixture_8(hass):
+    """Set up the mail_and_packages integration."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        title="imap.test.email",
+        data=FAKE_CONFIG_DATA_CAPOST,
+        version=CONFIG_VER,
+    )
+    entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    return entry
+
+
 @pytest.fixture()
 def mock_imap():
     """Mock imap class values."""
@@ -1740,3 +1757,30 @@ def mock_imap_amazon_otp():
         mock_conn.fetch.return_value = ("OK", [(b"", email_file.encode("utf-8"))])
         mock_conn.select.return_value = ("OK", [])
         yield mock_conn
+
+
+@pytest.fixture()
+def mock_imap_capost_mail():
+    """Mock imap class values."""
+    with patch(
+        "custom_components.mail_and_packages.helpers.imaplib"
+    ) as mock_imap_capost_mail:
+        mock_conn = mock.Mock(autospec=imaplib.IMAP4_SSL)
+        mock_imap_capost_mail.IMAP4_SSL.return_value = mock_conn
+
+        mock_conn.login.return_value = (
+            "OK",
+            [b"[email protected] authenticated (Success)"],
+        )
+        mock_conn.list.return_value = (
+            "OK",
+            [b'(\\HasNoChildren) "/" "INBOX"'],
+        )
+        mock_conn.search.return_value = ("OK", [b"1"])
+        mock_conn.uid.return_value = ("OK", [b"1"])
+        f = open("tests/test_emails/capost_mail.eml", "r")
+        email_file = f.read()
+        mock_conn.fetch.return_value = ("OK", [(b"", email_file.encode("utf-8"))])
+        mock_conn.select.return_value = ("OK", [])
+
+        yield mock_conn

tests/const.py

@@ -924,3 +924,28 @@
     "username": "[email protected]",
     "verify_ssl": False,
 }
+
+FAKE_CONFIG_DATA_CAPOST = {
+    "allow_external": False,
+    "custom_img": False,
+    "folder": '"INBOX"',
+    "generate_mp4": False,
+    "gif_duration": 5,
+    "host": "imap.test.email",
+    "image_name": "mail_today.gif",
+    "image_path": "custom_components/mail_and_packages/images/",
+    "image_security": True,
+    "imap_security": "SSL",
+    "imap_timeout": 30,
+    "password": "suchfakemuchpassword",
+    "port": 993,
+    "resources": [
+        "zpackages_delivered",
+        "zpackages_transit",
+        "capost_mail",
+    ],
+    "scan_interval": 20,
+    "storage": ".storage/mail_and_packages/images/",
+    "username": "[email protected]",
+    "verify_ssl": False,
+}