6 files changed, 36 insertions, 18 deletions

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index da22ae1..fb69ded 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -7,6 +7,7 @@ v1.6.3
 * Add server rules API (`instance_rules`)
 * Add confirmation email resend API (`email_resend_confirmation`)
 * Add account lookup API (`account_lookup`)
+* Add `policy` param to control notification sources for `push_subscription_set`
 
 v1.6.2
 ------
diff --git a/TODO.md b/TODO.md
index 6964ea3..bf73f06 100644
--- a/TODO.md
+++ b/TODO.md
@@ -32,7 +32,7 @@ Refer to mastodon changelog and API docs for details when implementing, add or m
 * [x] Add server rules

 * [x] Add POST /api/v1/emails/confirmations to REST API

 * [x] Add GET /api/v1/accounts/lookup to REST API

-* [ ] Add policy param to POST /api/v1/push/subscriptions in REST API

+* [x] Add policy param to POST /api/v1/push/subscriptions in REST API

 * [ ] Add details to error response for POST /api/v1/accounts in REST API

 

 3.4.2

diff --git a/mastodon/Mastodon.py b/mastodon/Mastodon.py
index 6e5bc29..12a054a 100644
--- a/mastodon/Mastodon.py
+++ b/mastodon/Mastodon.py
@@ -2756,7 +2756,7 @@ class Mastodon:
     def push_subscription_set(self, endpoint, encrypt_params, follow_events=None,
                               favourite_events=None, reblog_events=None,
                               mention_events=None, poll_events=None,
-                              follow_request_events=None):
+                              follow_request_events=None, status_events=None, policy='all'):
         """
         Sets up or modifies the push subscription the logged-in user has for this app.
 
@@ -2766,10 +2766,16 @@ class Mastodon:
         You can generate this as well as the corresponding private key using the
         `push_subscription_generate_keys()`_ function.
 
+        `policy` controls what sources will generate webpush events. Valid values are
+        `all`, `none`, `follower` and `followed`.
+
         The rest of the parameters controls what kind of events you wish to subscribe to.
 
         Returns a `push subscription dict`_.
         """
+        if not policy in ['all', 'none', 'follower', 'followed']:
+            raise MastodonIllegalArgumentError("Valid values for policy are 'all', 'none', 'follower' or 'followed'.")
+            
         endpoint = Mastodon.__protocolize(endpoint)
 
         push_pubkey_b64 = base64.b64encode(encrypt_params['pubkey'])
@@ -2778,7 +2784,8 @@ class Mastodon:
         params = {
             'subscription[endpoint]': endpoint,
             'subscription[keys][p256dh]': push_pubkey_b64,
-            'subscription[keys][auth]': push_auth_b64
+            'subscription[keys][auth]': push_auth_b64,
+            'policy': policy
         }
 
         if follow_events != None:
@@ -2799,6 +2806,9 @@ class Mastodon:
         if follow_request_events != None:
             params['data[alerts][follow_request]'] = follow_request_events
 
+        if follow_request_events != None:
+            params['data[alerts][status]'] = status_events
+
         # Canonicalize booleans
         params = self.__generate_params(params)
 
diff --git a/tests/cassettes/test_push_set.yaml b/tests/cassettes/test_push_set.yaml
index dd3089e..4995277 100644
--- a/tests/cassettes/test_push_set.yaml
+++ b/tests/cassettes/test_push_set.yaml
@@ -1,6 +1,6 @@
 interactions:
 - request:
-    body: subscription%5Bendpoint%5D=https%3A%2F%2Fexample.com&subscription%5Bkeys%5D%5Bp256dh%5D=BMGJORRZ33ebUVNTfhh5ONpJiYo1Qq6k%2FIv21xcIFzbR25IC2q7AE7RLc4x%2Ba5Rix8nZtyJ2QqVuBj5ScKvZvGQ%3D&subscription%5Bkeys%5D%5Bauth%5D=y6Jq8wK%2F5bnDIGR40sENrw%3D%3D
+    body: subscription%5Bendpoint%5D=https%3A%2F%2Fexample.com&subscription%5Bkeys%5D%5Bp256dh%5D=BKFO5w6Uf%2B%2F2wo89ovbphk5Zrb0mcAKjZrIyrX66f2IAijRtuXx4yK6J9hR%2FemKnF2DyQcyx7%2F4IGhKHBk0OTEk%3D&subscription%5Bkeys%5D%5Bauth%5D=6T8gVmd01DhQUbejRj%2Bxmg%3D%3D&policy=none&data%5Balerts%5D%5Bfollow%5D=1&data%5Balerts%5D%5Bfavourite%5D=1&data%5Balerts%5D%5Breblog%5D=1&data%5Balerts%5D%5Bmention%5D=1&data%5Balerts%5D%5Bpoll%5D=1&data%5Balerts%5D%5Bfollow_request%5D=1&data%5Balerts%5D%5Bstatus%5D=1
     headers:
       Accept:
       - '*/*'
@@ -11,7 +11,7 @@ interactions:
       Connection:
       - keep-alive
       Content-Length:
-      - '246'
+      - '489'
       Content-Type:
       - application/x-www-form-urlencoded
       User-Agent:
@@ -20,14 +20,14 @@ interactions:
     uri: http://localhost:3000/api/v1/push/subscription
   response:
     body:
-      string: '{"id":1,"endpoint":"https://example.com","alerts":{},"server_key":"BIX7IY8wYrdPf5wfG59B5B4CBbpQWmOBr_SI9zEsCDfMGDx8KmPuu-3WP7b-amaekv2NZj3ZSNst9OaqoxXvC1s="}'
+      string: '{"id":5,"endpoint":"https://example.com","alerts":{"mention":true,"status":true,"reblog":true,"follow":true,"follow_request":true,"favourite":true,"poll":true},"server_key":"BFu6DBpfcm8_h8gm3rHUkfaOLg7azvYN_auFI4KcNuh5SLBVMhTkKKvUaLENtA_c6v5Hmrucvh0WwsN1o9NFQRU="}'
     headers:
       Cache-Control:
       - no-store
       Content-Security-Policy:
       - 'base-uri ''none''; default-src ''none''; frame-ancestors ''none''; font-src
         ''self'' http://localhost:3000; img-src ''self'' https: data: blob: http://localhost:3000;
-        style-src ''self'' http://localhost:3000 ''nonce-IWpo+b1yj3qq43HAmklp8w=='';
+        style-src ''self'' http://localhost:3000 ''nonce-bz+9uelNbajqElylgkM2Gg=='';
         media-src ''self'' https: data: http://localhost:3000; frame-src ''self''
         https:; manifest-src ''self'' http://localhost:3000; connect-src ''self''
         data: blob: http://localhost:3000 http://localhost:3000 ws://localhost:4000
@@ -37,7 +37,7 @@ interactions:
       Content-Type:
       - application/json; charset=utf-8
       ETag:
-      - W/"07d4ac937919cd871dc7bf1552440cdd"
+      - W/"083c3b807a4d145ec452ffd03c768d76"
       Referrer-Policy:
       - strict-origin-when-cross-origin
       Transfer-Encoding:
@@ -53,9 +53,9 @@ interactions:
       X-Permitted-Cross-Domain-Policies:
       - none
       X-Request-Id:
-      - 095b046e-1a89-4d2c-8d0d-9af1c13fee60
+      - abb4aec7-a93a-4a99-ba4c-5a290bbbf782
       X-Runtime:
-      - '0.019991'
+      - '0.027761'
       X-XSS-Protection:
       - 1; mode=block
     status:
@@ -78,14 +78,14 @@ interactions:
     uri: http://localhost:3000/api/v1/push/subscription
   response:
     body:
-      string: '{"id":1,"endpoint":"https://example.com","alerts":{},"server_key":"BIX7IY8wYrdPf5wfG59B5B4CBbpQWmOBr_SI9zEsCDfMGDx8KmPuu-3WP7b-amaekv2NZj3ZSNst9OaqoxXvC1s="}'
+      string: '{"id":5,"endpoint":"https://example.com","alerts":{"mention":true,"status":true,"reblog":true,"follow":true,"follow_request":true,"favourite":true,"poll":true},"server_key":"BFu6DBpfcm8_h8gm3rHUkfaOLg7azvYN_auFI4KcNuh5SLBVMhTkKKvUaLENtA_c6v5Hmrucvh0WwsN1o9NFQRU="}'
     headers:
       Cache-Control:
       - no-store
       Content-Security-Policy:
       - 'base-uri ''none''; default-src ''none''; frame-ancestors ''none''; font-src
         ''self'' http://localhost:3000; img-src ''self'' https: data: blob: http://localhost:3000;
-        style-src ''self'' http://localhost:3000 ''nonce-F8ntBeTNhe2yPZYi1hFWAQ=='';
+        style-src ''self'' http://localhost:3000 ''nonce-4kBQ0XWCASsxXi+/Z0W5jA=='';
         media-src ''self'' https: data: http://localhost:3000; frame-src ''self''
         https:; manifest-src ''self'' http://localhost:3000; connect-src ''self''
         data: blob: http://localhost:3000 http://localhost:3000 ws://localhost:4000
@@ -95,7 +95,7 @@ interactions:
       Content-Type:
       - application/json; charset=utf-8
       ETag:
-      - W/"07d4ac937919cd871dc7bf1552440cdd"
+      - W/"083c3b807a4d145ec452ffd03c768d76"
       Referrer-Policy:
       - strict-origin-when-cross-origin
       Transfer-Encoding:
@@ -111,9 +111,9 @@ interactions:
       X-Permitted-Cross-Domain-Policies:
       - none
       X-Request-Id:
-      - 1fc1f34c-9edb-4647-a1b9-935865e14630
+      - eeb007da-a24b-4e31-b962-30ef0ba3bc16
       X-Runtime:
-      - '0.007914'
+      - '0.008203'
       X-XSS-Protection:
       - 1; mode=block
     status:
diff --git a/tests/test_instance.py b/tests/test_instance.py
index 5f451c7..8f3f142 100644
--- a/tests/test_instance.py
+++ b/tests/test_instance.py
@@ -42,7 +42,7 @@ def test_health(api):
 
 @pytest.mark.vcr()
 def test_server_time(api):
-     # present date...
+    # present date...
     present_time = api.get_approx_server_time()
     # hahahahaha
 
diff --git a/tests/test_push.py b/tests/test_push.py
index 0479191..daa99c7 100644
--- a/tests/test_push.py
+++ b/tests/test_push.py
@@ -26,11 +26,18 @@ def test_decrypt(api):
 @pytest.mark.vcr(match_on=['path'])
 def test_push_set(api):
     priv, pub = api.push_subscription_generate_keys()
-    sub = api.push_subscription_set("example.com", pub)
-    
+    sub = api.push_subscription_set("example.com", pub, follow_events=True, favourite_events=True, reblog_events=True, mention_events=True, poll_events=True, follow_request_events=True, status_events=True, policy='none')
     assert sub == api.push_subscription()
     assert sub.endpoint == "https://example.com"
 
+    should_throw = False
+    try:
+        sub = api.push_subscription_set("example.com", pub, follow_events=True, favourite_events=True, reblog_events=True, mention_events=True, poll_events=True, follow_request_events=True, status_events=True, policy='not a valid value')
+        should_throw = True
+    except:
+        pass
+    assert not should_throw
+
 @pytest.mark.vcr(match_on=['path'])
 def test_push_update(api):
     priv, pub = api.push_subscription_generate_keys()