diff options
| -rw-r--r-- | prometheus_pgbouncer_exporter/utils.py | 2 | ||||
| -rw-r--r-- | tests/test_connection.py | 12 |
2 files changed, 13 insertions, 1 deletions
diff --git a/prometheus_pgbouncer_exporter/utils.py b/prometheus_pgbouncer_exporter/utils.py index 3a99695..ea69f46 100644 --- a/prometheus_pgbouncer_exporter/utils.py +++ b/prometheus_pgbouncer_exporter/utils.py @@ -16,7 +16,7 @@ import psycopg2 def get_connection(user=None, port=None, host=None, dbname='pgbouncer', password=None): kwargs = { 'user': user, 'port': port, 'host': host, 'dbname': dbname, 'password': password } - kwargs = dict([(k, v) for k, v in kwargs.iteritems() if v]) + kwargs = dict([(k, v) for k, v in kwargs.items() if v]) connection = psycopg2.connect( **kwargs ) diff --git a/tests/test_connection.py b/tests/test_connection.py new file mode 100644 index 0000000..1c0f4ea --- /dev/null +++ b/tests/test_connection.py @@ -0,0 +1,12 @@ +import unittest + +from unittest.mock import patch, Mock + +from prometheus_pgbouncer_exporter import utils + +class ConnectionTest(unittest.TestCase): + @patch('prometheus_pgbouncer_exporter.utils.psycopg2.connect') + def test_get_connection_with_password_works(self, connect): + conn = utils.get_connection(host='/tmp/', dbname='template1') + connect.assert_called_once_with(host='/tmp/', dbname='template1') + connect.return_value.set_session.assert_called_once_with(autocommit=True) |