From 84b36536060a9ba8795e636a76eb9a92fdb506e3 Mon Sep 17 00:00:00 2001
From: Christian Decker <decker.christian@gmail.com>
Date: Sat, 2 Jan 2021 14:24:15 +0100
Subject: [PATCH] pyln: Add command notification support for plugins

Changelog-Added: pyln-client: Plugin methods can now report progress or status via the `Request.notify` function
---
 contrib/pyln-client/pyln/client/plugin.py | 55 ++++++++++++++++++-----
 tests/test_plugin.py                      | 15 ++++---
 2 files changed, 53 insertions(+), 17 deletions(-)

diff --git a/contrib/pyln-client/pyln/client/plugin.py b/contrib/pyln-client/pyln/client/plugin.py
index 98d0059a7..f86c44ef8 100644
--- a/contrib/pyln-client/pyln/client/plugin.py
+++ b/contrib/pyln-client/pyln/client/plugin.py
@@ -141,6 +141,47 @@ class Request(dict):
     def _write_result(self, result: dict) -> None:
         self.plugin._write_locked(result)
 
+    def _notify(self, method: str, params: JSONType) -> None:
+        """Send a notification to the caller.
+
+        Can contain a variety of things, but is usually used to report
+        progress or command status.
+
+        """
+        self._write_result({
+            'jsonrpc': '2.0',
+            'params': params,
+            "method": method,
+        })
+
+    def notify(self, message: str, level: str = 'info') -> None:
+        """Send a message notification to the caller.
+        """
+        self._notify(
+            "message",
+            params={
+                'id': self.id,
+                'level': level,
+                'message': message,
+            }
+        )
+
+    def progress(self,
+                 progress: int,
+                 total: int,
+                 stage: Optional[int] = None,
+                 stage_total: Optional[int] = None
+                 ) -> None:
+        d: Dict[str, JSONType] = {
+            "id": self.id,
+            "num": progress,
+            "total": total,
+        }
+        if stage is not None and stage_total is not None:
+            d['stage'] = {"num": stage, "total": stage_total}
+
+        self._notify("progress", d)
+
 
 # If a hook call fails we need to coerce it into something the main daemon can
 # handle. Returning an error is not an option since we explicitly do not allow
@@ -639,22 +680,14 @@ class Plugin(object):
     def notify_message(self, request: Request, message: str,
                        level: str = 'info') -> None:
         """Send a notification message to sender of this request"""
-        self.notify("message", {"id": request.id,
-                                "level": level,
-                                "message": message})
+        request.notify(message=message)
 
     def notify_progress(self, request: Request,
                         progress: int, progress_total: int,
                         stage: Optional[int] = None,
                         stage_total: Optional[int] = None) -> None:
-        """Send a progerss message to sender of this request: if more than one stage, set stage and stage_total"""
-        d: Dict[str, Any] = {"id": request.id,
-                             "num": progress,
-                             "total": progress_total}
-        if stage_total is not None:
-            d['stage'] = {"num": stage, "total": stage_total}
-
-        self.notify("progress", d)
+        """Send a progress message to sender of this request: if more than one stage, set stage and stage_total"""
+        request.progress(progress, progress_total, stage, stage_total)
 
     def _parse_request(self, jsrequest: Dict[str, JSONType]) -> Request:
         i = jsrequest.get('id', None)
diff --git a/tests/test_plugin.py b/tests/test_plugin.py
index 0953ce47f..ca4c505de 100644
--- a/tests/test_plugin.py
+++ b/tests/test_plugin.py
@@ -2027,15 +2027,18 @@ def test_notify(node_factory):
             assert out[2 + i].endswith("|\n")
         else:
             assert out[2 + i].endswith("|\r")
-    assert out[102] == '\r'
+
+    assert out[102] == '# Beginning stage 2\n'
+    assert out[103] == '\r'
+
     for i in range(10):
-        assert out[103 + i].startswith("# Stage 2/2 {:>2}/10 |".format(1 + i))
+        assert out[104 + i].startswith("# Stage 2/2 {:>2}/10 |".format(1 + i))
         if i == 9:
-            assert out[103 + i].endswith("|\n")
+            assert out[104 + i].endswith("|\n")
         else:
-            assert out[103 + i].endswith("|\r")
-    assert out[113] == '"This worked"\n'
-    assert len(out) == 114
+            assert out[104 + i].endswith("|\r")
+    assert out[114] == '"This worked"\n'
+    assert len(out) == 115
 
     # At debug level, we get the second prompt.
     out = subprocess.check_output(['cli/lightning-cli',