Fix tests

tests/test_main.py

@@ -14,6 +14,7 @@ from pgcli.main import (
     obfuscate_process_password,
     duration_in_words,
     format_output,
+    notify_callback,
     PGCli,
     OutputSettings,
     COLOR_CODE_REGEX,
@@ -433,6 +434,7 @@ def test_pg_service_file(tmpdir):
         "b_host",
         "5435",
         "",
+        notify_callback,
         application_name="pgcli",
     )
     del os.environ["PGPASSWORD"]
@@ -488,7 +490,7 @@ def test_application_name_db_uri(tmpdir):
         cli = PGCli(pgclirc_file=str(tmpdir.join("rcfile")))
         cli.connect_uri("postgres://bar@baz.com/?application_name=cow")
     mock_pgexecute.assert_called_with(
-        "bar", "bar", "", "baz.com", "", "", application_name="cow"
+        "bar", "bar", "", "baz.com", "", "", notify_callback, application_name="cow"
     )
 
 

tests/test_ssh_tunnel.py

@@ -6,7 +6,7 @@ from configobj import ConfigObj
 from click.testing import CliRunner
 from sshtunnel import SSHTunnelForwarder
 
-from pgcli.main import cli, PGCli
+from pgcli.main import cli, notify_callback, PGCli
 from pgcli.pgexecute import PGExecute
 
 
@@ -61,6 +61,7 @@ def test_ssh_tunnel(
         "127.0.0.1",
         pgcli.ssh_tunnel.local_bind_ports[0],
         "",
+        notify_callback,
     )
     mock_ssh_tunnel_forwarder.reset_mock()
     mock_pgexecute.reset_mock()
@@ -96,6 +97,7 @@ def test_ssh_tunnel(
         "127.0.0.1",
         pgcli.ssh_tunnel.local_bind_ports[0],
         "",
+        notify_callback,
     )
     mock_ssh_tunnel_forwarder.reset_mock()
     mock_pgexecute.reset_mock()