Diamond/src/collectors/postgres/test/testpostgres.py at b6eb57d1a69d47c678ecd17bd66127cd8e00333e · python-diamond/Diamond · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
#!/usr/bin/python
# coding=utf-8
##########################################################################

from test import CollectorTestCase
from test import get_collector_config

from mock import patch, Mock
from postgres import PostgresqlCollector


class TestPostgresqlCollector(CollectorTestCase):

    def setUp(self, allowed_names=None):
        if not allowed_names:
            allowed_names = []

        default_config = get_collector_config('PostgresqlCollector', {})
        self.default_collector = PostgresqlCollector(default_config, None)

        config = get_collector_config('PostgresqlCollector', {
            'password': 'default_password',
            'port': 5432,
            'instances': {
                'postgres_a': {
                    'host': 'db1.loc',
                    'password_provider': 'pgpass',
                },
                'postgres_b': {
                    'host': 'db2.loc',
                    'port': 5433,
                    'password': 'instance_password_b',
                }
            }
        })
        self.collector = PostgresqlCollector(config, None)

    def test_import(self):
        self.assertTrue(PostgresqlCollector)

    def test_config_override(self):
        self.assertEqual(self.collector._get_config('postgres_a', 'port'), 5432)

        self.assertEqual(self.collector._get_config('postgres_b', 'port'), 5433)

    @patch('postgres.psycopg2')
    def test_connect_with_password(self, psycopg2_mock):
        conn_mock = Mock()
        psycopg2_mock.connect.return_value = conn_mock

        ret = self.collector._connect('postgres_b', 'test_db')

        self.assertTrue(conn_mock.set_isolation_level.called)
        self.assertEqual(ret, conn_mock)
        psycopg2_mock.connect.assert_called_once_with(
            database='test_db', host='db2.loc', password='instance_password_b',
            port=5433, sslmode='disable', user='postgres'
        )

    @patch('postgres.psycopg2')
    def test_connect_with_pgpass(self, psycopg2_mock):
        conn_mock = Mock()
        psycopg2_mock.connect.return_value = conn_mock

        ret = self.collector._connect('postgres_a', 'test_db')

        self.assertTrue(conn_mock.set_isolation_level.called)
        self.assertEqual(ret, conn_mock)
        psycopg2_mock.connect.assert_called_once_with(
            database='test_db', host='db1.loc',
            port=5432, sslmode='disable', user='postgres'
        )

    @patch('postgres.psycopg2')
    def test_connect_error(self, psycopg2_mock):
        psycopg2_mock.connect.side_effect = Exception('Some db exc')

        with self.assertRaises(Exception):
            self.collector._connect('test_db')