Add SubscribeToAllFolders support to subscriptions (#1244)

ecederstrand · web-flow · commit e0543881c203 · 2023-10-31T19:06:58.000+01:00
* feat: add SubscribeToAllFolders support to subscriptions

* feat: Add context managers
diff --git a/exchangelib/account.py b/exchangelib/account.py
@@ -54,6 +54,7 @@
     ToDoSearch,
     VoiceMail,
 )
+from .folders.collections import PullSubscription, PushSubscription, StreamingSubscription
 from .items import ALL_OCCURRENCES, AUTO_RESOLVE, HARD_DELETE, ID_ONLY, SAVE_ONLY, SEND_TO_NONE
 from .properties import EWSElement, Mailbox, SendingAs
 from .protocol import Protocol
@@ -73,6 +74,10 @@
     MoveItem,
     SendItem,
     SetUserOofSettings,
+    SubscribeToPull,
+    SubscribeToPush,
+    SubscribeToStreaming,
+    Unsubscribe,
     UpdateItem,
     UploadItems,
 )
@@ -742,6 +747,73 @@ def delegates(self):
         """Return a list of DelegateUser objects representing the delegates that are set on this account."""
         return list(GetDelegate(account=self).call(user_ids=None, include_permissions=True))
 
+    def subscribe_to_pull(self, event_types=None, watermark=None, timeout=60):
+        """Create a pull subscription.
+
+        :param event_types: List of event types to subscribe to. Possible values defined in SubscribeToPull.EVENT_TYPES
+        :param watermark: An event bookmark as returned by some sync services
+        :param timeout: Timeout of the subscription, in minutes. Timeout is reset when the server receives a
+        GetEvents request for this subscription.
+        :return: The subscription ID and a watermark
+        """
+        if event_types is None:
+            event_types = SubscribeToPull.EVENT_TYPES
+        return SubscribeToPull(account=self).get(
+            folders=None,
+            event_types=event_types,
+            watermark=watermark,
+            timeout=timeout,
+        )
+
+    def subscribe_to_push(self, callback_url, event_types=None, watermark=None, status_frequency=1):
+        """Create a push subscription.
+
+        :param callback_url: A client-defined URL that the server will call
+        :param event_types: List of event types to subscribe to. Possible values defined in SubscribeToPush.EVENT_TYPES
+        :param watermark: An event bookmark as returned by some sync services
+        :param status_frequency: The frequency, in minutes, that the callback URL will be called with.
+        :return: The subscription ID and a watermark
+        """
+        if event_types is None:
+            event_types = SubscribeToPush.EVENT_TYPES
+        return SubscribeToPush(account=self).get(
+            folders=None,
+            event_types=event_types,
+            watermark=watermark,
+            status_frequency=status_frequency,
+            url=callback_url,
+        )
+
+    def subscribe_to_streaming(self, event_types=None):
+        """Create a streaming subscription.
+
+        :param event_types: List of event types to subscribe to. Possible values defined in SubscribeToPush.EVENT_TYPES
+        :return: The subscription ID
+        """
+        if event_types is None:
+            event_types = SubscribeToStreaming.EVENT_TYPES
+        return SubscribeToStreaming(account=self).get(folders=None, event_types=event_types)
+
+    def pull_subscription(self, **kwargs):
+        return PullSubscription(target=self, **kwargs)
+
+    def push_subscription(self, **kwargs):
+        return PushSubscription(target=self, **kwargs)
+
+    def streaming_subscription(self, **kwargs):
+        return StreamingSubscription(target=self, **kwargs)
+
+    def unsubscribe(self, subscription_id):
+        """Unsubscribe. Only applies to pull and streaming notifications.
+
+        :param subscription_id: A subscription ID as acquired by .subscribe_to_[pull|streaming]()
+        :return: True
+
+        This method doesn't need the current collection instance, but it makes sense to keep the method along the other
+        sync methods.
+        """
+        return Unsubscribe(account=self).get(subscription_id=subscription_id)
+
     def __str__(self):
         if self.fullname:
             return f"{self.primary_smtp_address} ({self.fullname})"
diff --git a/exchangelib/folders/base.py b/exchangelib/folders/base.py
@@ -631,15 +631,15 @@ def subscribe_to_streaming(self, event_types=None):
 
     @require_id
     def pull_subscription(self, **kwargs):
-        return PullSubscription(folder=self, **kwargs)
+        return PullSubscription(target=self, **kwargs)
 
     @require_id
     def push_subscription(self, **kwargs):
-        return PushSubscription(folder=self, **kwargs)
+        return PushSubscription(target=self, **kwargs)
 
     @require_id
     def streaming_subscription(self, **kwargs):
-        return StreamingSubscription(folder=self, **kwargs)
+        return StreamingSubscription(target=self, **kwargs)
 
     def unsubscribe(self, subscription_id):
         """Unsubscribe. Only applies to pull and streaming notifications.
diff --git a/exchangelib/folders/collections.py b/exchangelib/folders/collections.py
@@ -448,13 +448,13 @@ def subscribe_to_streaming(self, event_types=None):
         return SubscribeToStreaming(account=self.account).get(folders=self.folders, event_types=event_types)
 
     def pull_subscription(self, **kwargs):
-        return PullSubscription(folder=self, **kwargs)
+        return PullSubscription(target=self, **kwargs)
 
     def push_subscription(self, **kwargs):
-        return PushSubscription(folder=self, **kwargs)
+        return PushSubscription(target=self, **kwargs)
 
     def streaming_subscription(self, **kwargs):
-        return StreamingSubscription(folder=self, **kwargs)
+        return StreamingSubscription(target=self, **kwargs)
 
     def unsubscribe(self, subscription_id):
         """Unsubscribe. Only applies to pull and streaming notifications.
@@ -540,8 +540,8 @@ def sync_hierarchy(self, sync_state=None, only_fields=None):
 
 
 class BaseSubscription(metaclass=abc.ABCMeta):
-    def __init__(self, folder, **subscription_kwargs):
-        self.folder = folder
+    def __init__(self, target, **subscription_kwargs):
+        self.target = target
         self.subscription_kwargs = subscription_kwargs
         self.subscription_id = None
 
@@ -550,19 +550,19 @@ def __enter__(self):
         """Create the subscription"""
 
     def __exit__(self, *args, **kwargs):
-        self.folder.unsubscribe(subscription_id=self.subscription_id)
+        self.target.unsubscribe(subscription_id=self.subscription_id)
         self.subscription_id = None
 
 
 class PullSubscription(BaseSubscription):
     def __enter__(self):
-        self.subscription_id, watermark = self.folder.subscribe_to_pull(**self.subscription_kwargs)
+        self.subscription_id, watermark = self.target.subscribe_to_pull(**self.subscription_kwargs)
         return self.subscription_id, watermark
 
 
 class PushSubscription(BaseSubscription):
     def __enter__(self):
-        self.subscription_id, watermark = self.folder.subscribe_to_push(**self.subscription_kwargs)
+        self.subscription_id, watermark = self.target.subscribe_to_push(**self.subscription_kwargs)
         return self.subscription_id, watermark
 
     def __exit__(self, *args, **kwargs):
@@ -572,5 +572,5 @@ def __exit__(self, *args, **kwargs):
 
 class StreamingSubscription(BaseSubscription):
     def __enter__(self):
-        self.subscription_id = self.folder.subscribe_to_streaming(**self.subscription_kwargs)
+        self.subscription_id = self.target.subscribe_to_streaming(**self.subscription_kwargs)
         return self.subscription_id
diff --git a/exchangelib/services/subscribe.py b/exchangelib/services/subscribe.py
@@ -40,9 +40,13 @@ def _get_elements_in_container(cls, container):
         return [(container.find(f"{{{MNS}}}SubscriptionId"), container.find(f"{{{MNS}}}Watermark"))]
 
     def _partial_payload(self, folders, event_types):
-        request_elem = create_element(self.subscription_request_elem_tag)
-        folder_ids = folder_ids_element(folders=folders, version=self.account.version, tag="t:FolderIds")
-        request_elem.append(folder_ids)
+        if folders is None:
+            # Interpret this as "all folders"
+            request_elem = create_element(self.subscription_request_elem_tag, attrs=dict(SubscribeToAllFolders=True))
+        else:
+            request_elem = create_element(self.subscription_request_elem_tag)
+            folder_ids = folder_ids_element(folders=folders, version=self.account.version, tag="t:FolderIds")
+            request_elem.append(folder_ids)
         event_types_elem = create_element("t:EventTypes")
         for event_type in event_types:
             add_xml_child(event_types_elem, "t:EventType", event_type)
diff --git a/tests/test_items/test_sync.py b/tests/test_items/test_sync.py
@@ -54,6 +54,26 @@ def test_pull_subscribe(self):
             self.account.root.tois.children.unsubscribe(subscription_id)
         # Affinity cookie is not always sent by the server for pull subscriptions
 
+    def test_pull_subscribe_from_account(self):
+        self.account.affinity_cookie = None
+        with self.account.pull_subscription() as (subscription_id, watermark):
+            self.assertIsNotNone(subscription_id)
+            self.assertIsNotNone(watermark)
+        # Test with watermark
+        with self.account.pull_subscription(watermark=watermark) as (subscription_id, watermark):
+            self.assertIsNotNone(subscription_id)
+            self.assertIsNotNone(watermark)
+        # Context manager already unsubscribed us
+        with self.assertRaises(ErrorSubscriptionNotFound):
+            self.account.unsubscribe(subscription_id)
+        # Test without watermark
+        with self.account.pull_subscription() as (subscription_id, watermark):
+            self.assertIsNotNone(subscription_id)
+            self.assertIsNotNone(watermark)
+        with self.assertRaises(ErrorSubscriptionNotFound):
+            self.account.unsubscribe(subscription_id)
+        # Affinity cookie is not always sent by the server for pull subscriptions
+
     def test_push_subscribe(self):
         with self.account.inbox.push_subscription(callback_url="https://example.com/foo") as (
             subscription_id,
@@ -81,6 +101,33 @@ def test_push_subscribe(self):
         with self.assertRaises(ErrorInvalidSubscription):
             self.account.root.tois.children.unsubscribe(subscription_id)
 
+    def test_push_subscribe_from_account(self):
+        with self.account.push_subscription(callback_url="https://example.com/foo") as (
+            subscription_id,
+            watermark,
+        ):
+            self.assertIsNotNone(subscription_id)
+            self.assertIsNotNone(watermark)
+        # Test with watermark
+        with self.account.push_subscription(
+            callback_url="https://example.com/foo",
+            watermark=watermark,
+        ) as (subscription_id, watermark):
+            self.assertIsNotNone(subscription_id)
+            self.assertIsNotNone(watermark)
+        # Cannot unsubscribe. Must be done as response to callback URL request
+        with self.assertRaises(ErrorInvalidSubscription):
+            self.account.unsubscribe(subscription_id)
+        # Test via folder collection
+        with self.account.push_subscription(callback_url="https://example.com/foo") as (
+            subscription_id,
+            watermark,
+        ):
+            self.assertIsNotNone(subscription_id)
+            self.assertIsNotNone(watermark)
+        with self.assertRaises(ErrorInvalidSubscription):
+            self.account.unsubscribe(subscription_id)
+
     def test_empty_folder_collection(self):
         self.assertEqual(FolderCollection(account=None, folders=[]).subscribe_to_pull(), None)
         self.assertEqual(FolderCollection(account=None, folders=[]).subscribe_to_push("http://example.com"), None)
@@ -102,6 +149,22 @@ def test_streaming_subscribe(self):
         # Test affinity cookie
         self.assertIsNotNone(self.account.affinity_cookie)
 
+    def test_streaming_subscribe_from_account(self):
+        self.account.affinity_cookie = None
+        with self.account.streaming_subscription() as subscription_id:
+            self.assertIsNotNone(subscription_id)
+        # Context manager already unsubscribed us
+        with self.assertRaises(ErrorSubscriptionNotFound):
+            self.account.unsubscribe(subscription_id)
+        # Test via folder collection
+        with self.account.streaming_subscription() as subscription_id:
+            self.assertIsNotNone(subscription_id)
+        with self.assertRaises(ErrorSubscriptionNotFound):
+            self.account.unsubscribe(subscription_id)
+
+        # Test affinity cookie
+        self.assertIsNotNone(self.account.affinity_cookie)
+
     def test_sync_folder_hierarchy(self):
         test_folder = self.get_test_folder().save()
 
