Add test

pgcli/main.py

@@ -128,6 +128,15 @@ class PgCliQuitError(Exception):
     pass
 
 
+def notify_callback(notify: Notify):
+    click.secho(
+        'Notification received on channel "{}" (PID {}):\n{}'.format(
+            notify.channel, notify.pid, notify.payload
+        ),
+        fg="green",
+    )
+
+
 class PGCli:
     default_prompt = "\\u@\\h:\\d> "
     max_len_prompt = 30
@@ -654,14 +663,6 @@ class PGCli:
             if dsn:
                 dsn = make_conninfo(dsn, host=host, port=port)
 
-        def notify_callback(notify: Notify):
-            click.secho(
-                'Notification received on channel "{}" (PID {}):\n{}'.format(
-                    notify.channel, notify.pid, notify.payload
-                ),
-                fg="green",
-            )
-
         # Attempt to connect to the database.
         # Note that passwd may be empty on the first attempt. If connection
         # fails because of a missing or incorrect password, but we're allowed to

tests/conftest.py

@@ -9,6 +9,7 @@ from utils import (
     db_connection,
     drop_tables,
 )
+import pgcli.main
 import pgcli.pgexecute
 
 
@@ -37,6 +38,7 @@ def executor(connection):
         password=POSTGRES_PASSWORD,
         port=POSTGRES_PORT,
         dsn=None,
+        notify_callback=pgcli.main.notify_callback,
     )
 
 

tests/test_main.py

@@ -1,5 +1,6 @@
 import os
 import platform
+import re
 from unittest import mock
 
 import pytest
@@ -514,3 +515,24 @@ def test_application_name_db_uri(tmpdir):
 )
 def test_duration_in_words(duration_in_seconds, words):
     assert duration_in_words(duration_in_seconds) == words
+
+
+@dbtest
+def test_notifications(executor):
+    run(executor, "listen chan1")
+
+    with mock.patch("pgcli.main.click.secho") as mock_secho:
+        run(executor, "notify chan1, 'testing1'")
+        mock_secho.assert_called()
+        arg = mock_secho.call_args_list[0].args[0]
+    assert re.match(
+        r'Notification received on channel "chan1" \(PID \d+\):\ntesting1',
+        arg,
+    )
+
+    run(executor, "unlisten chan1")
+
+    with mock.patch("pgcli.main.click.secho") as mock_secho:
+        run(executor, "notify chan1, 'testing2'")
+        mock_secho.assert_not_called()
+