From 109b8d8909b1acab3a5a55ed20733d379e652d72 Mon Sep 17 00:00:00 2001
From: Ryan Barrett <git@ryanb.org>
Date: Tue, 8 Aug 2023 20:34:36 -0700
Subject: [PATCH] flask_server: add subscription support

---
 lexrpc/flask_server.py            | 7 +++++++
 lexrpc/tests/test_flask_server.py | 8 ++++++++
 2 files changed, 15 insertions(+)

diff --git a/lexrpc/flask_server.py b/lexrpc/flask_server.py
index 8bc2443..e4ddf65 100644
--- a/lexrpc/flask_server.py
+++ b/lexrpc/flask_server.py
@@ -4,6 +4,7 @@
 from flask import request
 from flask.json import jsonify
 from flask.views import View
+from flask_sock import Sock
 from jsonschema import ValidationError
 
 from .base import NSID_RE
@@ -26,6 +27,12 @@ def init_flask(xrpc_server, app):
       app: :class:`flask.Flask`
     """
     logger.info(f'Registering {xrpc_server} with {app}')
+
+    sock = Sock(app)
+    for nsid, fn in xrpc_server._methods.items():
+        if xrpc_server._defs[nsid]['type'] == 'subscription':
+            sock.route(f'/xrpc/{nsid}')(fn)
+
     app.add_url_rule('/xrpc/<nsid>',
                      view_func=XrpcEndpoint.as_view('xrpc-endpoint', xrpc_server),
                      methods=['GET', 'POST', 'OPTIONS'])
diff --git a/lexrpc/tests/test_flask_server.py b/lexrpc/tests/test_flask_server.py
index 997aa21..2fb047f 100644
--- a/lexrpc/tests/test_flask_server.py
+++ b/lexrpc/tests/test_flask_server.py
@@ -63,6 +63,14 @@ def test_query_boolean_param(self):
         self.assertEqual("Got 'foolz' for boolean parameter z, expected true or false",
                          resp.json['message'])
 
+    # TODO
+    # needs websocket test client, but flask-sock doesn't have one yet
+    # https://github.com/miguelgrinberg/flask-sock/issues/23
+    @skip
+    def test_subscription(self):
+        resp = self.client.get('/xrpc/io.example.subscribe?start=3&end=6')
+        self.assertEqual(200, resp.status_code, resp.json)
+
     # TODO
     @skip
     def test_procedure_missing_input(self):