From 9d559cc41b979d92aa10cf4c6cb49f3a1f00a06f Mon Sep 17 00:00:00 2001 From: Irina Truong Date: Fri, 8 Mar 2019 11:20:01 -0800 Subject: Unit test. --- tests/test_pgexecute.py | 34 ++++++++++++++++++++++++++++++---- 1 file changed, 30 insertions(+), 4 deletions(-) diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py index 39788950..8501837f 100644 --- a/tests/test_pgexecute.py +++ b/tests/test_pgexecute.py @@ -1,13 +1,16 @@ # coding=UTF-8 from __future__ import print_function -import pytest -import psycopg2 -from mock import patch -from pgcli.packages.parseutils.meta import FunctionMetadata from textwrap import dedent + +import psycopg2 +import pytest +from mock import patch, MagicMock +from pgspecial.main import PGSpecial, NO_QUERY from utils import run, dbtest, requires_json, requires_jsonb + from pgcli.main import PGCli +from pgcli.packages.parseutils.meta import FunctionMetadata def function_meta_data( @@ -425,3 +428,26 @@ def test_short_host(executor): assert executor.short_host == 'localhost' with patch.object(executor, 'host', 'localhost1.example.org,localhost2.example.org'): assert executor.short_host == 'localhost1' + + +class BrokenConnection(object): + """Mock a connection that failed.""" + + def cursor(self): + raise psycopg2.InterfaceError("I'm broken!") + + +@dbtest +def test_exit_without_active_connection(executor): + quit_handler = MagicMock() + pgspecial = PGSpecial() + pgspecial.register(quit_handler, '\\q', '\\q', 'Quit pgcli.', arg_type=NO_QUERY, case_sensitive=True, aliases=(':q',)) + + with patch.object(executor, "conn", BrokenConnection()): + # we should be able to quit the app, even without active connection + run(executor, "\\q", pgspecial=pgspecial) + quit_handler.assert_called_once() + + # an exception should be raised when running a query without active connection + with pytest.raises(psycopg2.InterfaceError): + run(executor, "select 1", pgspecial=pgspecial) -- cgit v1.2.3