From 3c55d7c145a38bd265b1e47f4082a3f61d83ff2a Mon Sep 17 00:00:00 2001
From: Ryan Barrett <git@ryanb.org>
Date: Fri, 19 Apr 2024 12:11:08 -0700
Subject: [PATCH] protocol: extract out enable/disable_protocol methods

for #880
---
 models.py   | 28 +++++++++++++++++++++++++++-
 protocol.py | 18 +++---------------
 2 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/models.py b/models.py
index f27b6f6..f72db3d 100644
--- a/models.py
+++ b/models.py
@@ -28,7 +28,7 @@ from oauth_dropins.webutil.models import (
 from oauth_dropins.webutil.util import json_dumps, json_loads
 
 import common
-from common import add, base64_to_long, DOMAIN_RE, long_to_base64, unwrap
+from common import add, base64_to_long, DOMAIN_RE, long_to_base64, remove, unwrap
 import ids
 
 # maps string label to Protocol subclass. populated by ProtocolUserMeta.
@@ -348,6 +348,32 @@ class User(StringIdModel, metaclass=ProtocolUserMeta):
 
         return None
 
+    @ndb.transactional()
+    def enable_protocol(self, to_proto):
+        """Adds ``to_proto` to :attr:`enabled_protocols`.
+
+        Args:
+          to_proto (:class:`protocol.Protocol` subclass)
+        """
+        user = self.key.get()
+        add(user.enabled_protocols, to_proto.LABEL)
+        user.put()
+
+        add(self.enabled_protocols, to_proto.LABEL)
+
+    @ndb.transactional()
+    def disable_protocol(self, to_proto):
+        """Removes ``to_proto` from :attr:`enabled_protocols`.
+
+        Args:
+          to_proto (:class:`protocol.Protocol` subclass)
+        """
+        user = self.key.get()
+        remove(user.enabled_protocols, to_proto.LABEL)
+        user.put()
+
+        remove(self.enabled_protocols, to_proto.LABEL)
+
     def handle_as(self, to_proto):
         """Returns this user's handle in a different protocol.
 
diff --git a/protocol.py b/protocol.py
index 1ea86d6..0659007 100644
--- a/protocol.py
+++ b/protocol.py
@@ -27,7 +27,6 @@ from common import (
     DOMAINS,
     error,
     PROTOCOL_DOMAINS,
-    remove,
     subdomain_wrap,
 )
 from flask_app import app
@@ -802,13 +801,7 @@ class Protocol:
                 logger.info("Ignoring block, target isn't one of our protocol domains")
                 return 'OK', 200
 
-            @ndb.transactional()
-            def disable_protocol():
-                user = from_user.key.get()
-                remove(user.enabled_protocols, proto.LABEL)
-                user.put()
-
-            disable_protocol()
+            from_user.disable_protocol(proto)
             return 'OK', 200
 
         # fetch actor if necessary
@@ -837,13 +830,8 @@ class Protocol:
             proto = Protocol.for_bridgy_subdomain(inner_obj_id)
             if proto:
                 # follow of one of our protocol users; enable that protocol
-                @ndb.transactional()
-                def enable_protocol():
-                    user = from_user.key.get()
-                    add(user.enabled_protocols, proto.LABEL)
-                    user.put()
-
-                enable_protocol()
+                from_user.enable_protocol(proto)
+                # TODO: accept
                 return 'OK', 200
 
             from_cls.handle_follow(obj)