310 lines
12 KiB
Python
310 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
import unittest
|
|
import psycopg2
|
|
import time
|
|
from pgbouncemgr.postgres import *
|
|
from pgbouncemgr.logger import Logger
|
|
from pgbouncemgr.node_config import NodeConfig
|
|
from tests.stub_psycopg2 import *
|
|
|
|
|
|
class PgConnectionTests(unittest.TestCase):
|
|
def test_GivenAuthenticationFailure_PgConnection_RaisesException(self):
|
|
self._test_connect_exception(
|
|
PgConnection, StubPsycopg2().add_auth_failure(),
|
|
PgConnectionFailed, "authentication failed")
|
|
|
|
def test_GivenConnectionTimeout_PgConnection_RaisesException(self):
|
|
self._test_connect_exception(
|
|
PgConnection, StubPsycopg2().add_conn_timeout(),
|
|
PgConnectionFailed, "timeout expired")
|
|
|
|
def test_GivenConnectionFailure_PgConnection_RaisesException(self):
|
|
self._test_connect_exception(
|
|
PgConnection, StubPsycopg2().add_conn_failure(),
|
|
PgConnectionFailed, "could not connect")
|
|
|
|
def test_GivenPostgresStartingUp_PgConnection_RaisesException(self):
|
|
self._test_connect_exception(
|
|
PgConnection, StubPsycopg2().add_admin_startup(),
|
|
PgConnectionFailed, "system is starting up")
|
|
|
|
def test_GivenPostgresShuttingDown_PgConnection_RaisesException(self):
|
|
self._test_connect_exception(
|
|
PgConnection, StubPsycopg2().add_admin_shutdown(),
|
|
PgConnectionFailed, "AdminShutdown")
|
|
|
|
def _test_connect_exception(self, test_class, stub_psycopg2, err, msg):
|
|
pg = test_class(NodeConfig(1), Logger(), stub_psycopg2)
|
|
with self.assertRaises(err) as context:
|
|
pg.connect()
|
|
self.assertIn(msg, str(context.exception))
|
|
|
|
def test_NodeConfig_IsAppliedToConnParams(self):
|
|
node_config = NodeConfig(1)
|
|
node_config.host = "1.1.1.1"
|
|
node_config.port = 9999
|
|
pg = PgConnection(node_config, Logger(), None)
|
|
|
|
self.assertEqual(1, pg.node_id)
|
|
self.assertEqual("1.1.1.1", pg.conn_params["host"])
|
|
self.assertEqual(9999, pg.conn_params["port"])
|
|
|
|
def test_GivenNoneValueInNodeConfig_ValueIsOmittedInConnParams(self):
|
|
node_config = NodeConfig(2)
|
|
node_config.host = None
|
|
node_config.port = None
|
|
stub_psycopg2 = StubPsycopg2().add_connection(StubConnection())
|
|
pg = PgConnection(node_config, Logger(), stub_psycopg2)
|
|
pg.connect()
|
|
|
|
self.assertEqual(2, pg.node_id)
|
|
self.assertNotIn("host", pg.conn_params)
|
|
self.assertNotIn("port", pg.conn_params)
|
|
|
|
def test_FirstConnect_SetsUpConnection(self):
|
|
stub_psycopg2 = StubPsycopg2().add_connection(StubConnection())
|
|
pg = PgConnection(NodeConfig('a'), Logger(), stub_psycopg2)
|
|
|
|
result = pg.connect()
|
|
|
|
self.assertEqual("CONN_CONNECTED", result)
|
|
|
|
def test_SecondConnect_PingsAndReusesConnection(self):
|
|
cursor = StubCursor(("SELECT", [[1]]))
|
|
conn = StubConnection(cursor)
|
|
stub_psycopg2 = StubPsycopg2().add_connection(conn)
|
|
pg = PgConnection(NodeConfig('b'), Logger(), stub_psycopg2)
|
|
|
|
result1 = pg.connect()
|
|
self.assertIs(None, cursor.query)
|
|
result2 = pg.connect()
|
|
self.assertEqual("SELECT 1", cursor.query)
|
|
self.assertEqual("CONN_CONNECTED", result1)
|
|
self.assertEqual("CONN_REUSED", result2)
|
|
|
|
def test_SecondConnectPingFails_SetsUpNewConnection(self):
|
|
conn1 = StubConnection(StubCursor(psycopg2.OperationalError()))
|
|
conn2 = StubConnection()
|
|
stub_psycopg2 = StubPsycopg2(conn1, conn2)
|
|
pg = PgConnection(NodeConfig('b'), Logger(), stub_psycopg2)
|
|
|
|
result1 = pg.connect() # Connection OK
|
|
result2 = pg.connect() # Ping fails, reconnectt
|
|
|
|
self.assertEqual("CONN_CONNECTED", result1)
|
|
self.assertEqual("CONN_RECONNECTED", result2)
|
|
|
|
def test_Disconnect_ClosesConnection(self):
|
|
conn = StubConnection()
|
|
stub_psycopg2 = StubPsycopg2(conn)
|
|
pg = PgConnection(NodeConfig('disco'), Logger(), stub_psycopg2)
|
|
|
|
pg.connect()
|
|
self.assertTrue(conn.connected)
|
|
pg.disconnect()
|
|
self.assertFalse(conn.connected)
|
|
|
|
def test_BigConnectionFlow(self):
|
|
conn1 = StubConnection(
|
|
StubCursor(psycopg2.OperationalError()))
|
|
conn2 = StubConnection(
|
|
StubCursor(("SELECT", [])),
|
|
StubCursor(("SELECT", [])))
|
|
conn3 = StubConnection(
|
|
StubCursor(("SELECT", [])))
|
|
|
|
stub_psycopg2 = StubPsycopg2(conn1, conn2, conn3)
|
|
pg = PgConnection(NodeConfig('b'), Logger(), stub_psycopg2)
|
|
|
|
result1 = pg.connect() # Connection 1 OK
|
|
result2 = pg.connect() # Ping fails, new connection 2 OK
|
|
result3 = pg.connect() # Ping 1 success
|
|
result4 = pg.connect() # Ping 2 success
|
|
pg.disconnect() # Explicit disconnect
|
|
result5 = pg.connect() # Connection 3 OK
|
|
result6 = pg.connect() # Ping success
|
|
|
|
self.assertEqual("CONN_CONNECTED", result1)
|
|
self.assertEqual("CONN_RECONNECTED", result2)
|
|
self.assertEqual("CONN_REUSED", result3)
|
|
self.assertEqual("CONN_REUSED", result4)
|
|
self.assertEqual("CONN_CONNECTED", result5)
|
|
self.assertEqual("CONN_REUSED", result6)
|
|
|
|
|
|
class PgReplicationConnectionTests(unittest.TestCase):
|
|
def test_LogicalReplicationConnection_IsUsed(self):
|
|
conn = StubConnection()
|
|
stub_psycopg2 = StubPsycopg2().add_connection(conn)
|
|
pg = PgReplicationConnection(
|
|
NodeConfig("foo"), Logger(), stub_psycopg2)
|
|
|
|
pg.connect()
|
|
self.assertEqual(
|
|
psycopg2.extras.LogicalReplicationConnection,
|
|
pg.conn_params["connection_factory"])
|
|
|
|
def test_GivenFailingConnection_ReplicationStatusIsOffline(self):
|
|
stub_psycopg2 = StubPsycopg2().add_auth_failure()
|
|
pg = PgReplicationConnection(
|
|
NodeConfig("foo"), Logger(), stub_psycopg2)
|
|
|
|
status = pg.get_replication_status()
|
|
self.assertEqual({
|
|
"status": "NODE_OFFLINE",
|
|
"system_id": None,
|
|
"timeline_id": None}, status)
|
|
|
|
def test_GivenFailingStandbyQuery_ReplicationStatusRaisesException(self):
|
|
conn = StubConnection(StubCursor(psycopg2.OperationalError()))
|
|
pg = PgReplicationConnection(
|
|
NodeConfig("foo"), Logger(), StubPsycopg2(conn))
|
|
|
|
with self.assertRaises(RetrievingPgReplicationStatusFailed) as context:
|
|
pg.get_replication_status()
|
|
self.assertIn("pg_is_in_recovery() failed", str(context.exception))
|
|
|
|
def test_GivenFailingIdentifySystemQuery_ReplicationStatusRaisesException(self):
|
|
conn = StubConnection(
|
|
StubCursor(("SELECT", [[False]])),
|
|
psycopg2.OperationalError())
|
|
pg = PgReplicationConnection(
|
|
NodeConfig("foo"), Logger(), StubPsycopg2(conn))
|
|
|
|
with self.assertRaises(RetrievingPgReplicationStatusFailed) as context:
|
|
pg.get_replication_status()
|
|
self.assertIn("IDENTIFY_SYSTEM failed", str(context.exception))
|
|
|
|
def test_GivenConnectionToPrimaryNode_ReplicationStatusIsPrimary(self):
|
|
conn = StubConnection(
|
|
StubCursor(("SELECT", [[False]])),
|
|
StubCursor(("IDENTIFY_SYSTEM", [["id", 1234, "other", "fields"]])))
|
|
stub_psycopg2 = StubPsycopg2(conn)
|
|
pg = PgReplicationConnection(
|
|
NodeConfig("foo"), Logger(), stub_psycopg2)
|
|
|
|
status = pg.get_replication_status()
|
|
self.assertEqual({
|
|
"status": "NODE_PRIMARY",
|
|
"system_id": "id",
|
|
"timeline_id": 1234}, status)
|
|
|
|
def test_GivenConnectionToStandbyNode_ReplicationStatusIsStandby(self):
|
|
conn = StubConnection(
|
|
StubCursor(("SELECT", [[True]])),
|
|
StubCursor(("IDENTIFY_SYSTEM", [["towel", 42, "other", "fields"]])))
|
|
stub_psycopg2 = StubPsycopg2(conn)
|
|
pg = PgReplicationConnection(
|
|
NodeConfig("foo"), Logger(), stub_psycopg2)
|
|
|
|
status = pg.get_replication_status()
|
|
self.assertEqual({
|
|
"status": "NODE_STANDBY",
|
|
"system_id": "towel",
|
|
"timeline_id": 42}, status)
|
|
|
|
|
|
class PgConnectionViaPgbouncerTests(unittest.TestCase):
|
|
def test_NodeConfigAndPgBouncerConfig_AreMergedInConnParams(self):
|
|
node_config = NodeConfig(777)
|
|
node_config.host = "1.1.1.1"
|
|
node_config.port = 9999
|
|
pgbouncer_config = {"host": "192.168.0.1", "port": 7654}
|
|
|
|
pg = PgConnectionViaPgbouncer(
|
|
node_config, pgbouncer_config, Logger(), None)
|
|
|
|
self.assertEqual(777, pg.node_id)
|
|
self.assertEqual("template1", pg.conn_params["database"])
|
|
self.assertEqual("192.168.0.1", pg.conn_params["host"])
|
|
self.assertEqual(7654, pg.conn_params["port"])
|
|
|
|
def test_WhenVerifyQueryCannotConnect_ExceptionIsRaised(self):
|
|
stub_psycopg2 = StubPsycopg2().add_auth_failure()
|
|
node_config = NodeConfig('xyz123')
|
|
pgbouncer_config = {"host": "192.168.0.1", "port": 7654}
|
|
pg = PgConnectionViaPgbouncer(
|
|
node_config, pgbouncer_config, Logger(), stub_psycopg2)
|
|
|
|
ok, err, cursor = pg.verify_connection()
|
|
self.assertFalse(ok)
|
|
self.assertIs(None, cursor)
|
|
self.assertIn("password authentication failed", str(err))
|
|
|
|
def test_WhenVerifyQueryTimesOut_ExceptionIsRaised(self):
|
|
def sleeping_query(query):
|
|
time.sleep(1)
|
|
conn = StubConnection(StubCursor(sleeping_query))
|
|
stub_psycopg2 = StubPsycopg2(conn)
|
|
node_config = NodeConfig('xyz123')
|
|
node_config.connect_timeout = 0.01
|
|
pgbouncer_config = {"host": "192.168.0.1", "port": 7654}
|
|
pg = PgConnectionViaPgbouncer(
|
|
node_config, pgbouncer_config, Logger(), stub_psycopg2)
|
|
|
|
ok, err, cursor = pg.verify_connection()
|
|
self.assertFalse(ok)
|
|
self.assertIs(None, cursor)
|
|
self.assertIn("Connection attempt timed out", str(err))
|
|
|
|
def test_WhenVerifyQueryRuns_NodeHostAndPortAreQueried(self):
|
|
conn = StubConnection(StubCursor(('SELECT', [[None]])))
|
|
stub_psycopg2 = StubPsycopg2(conn)
|
|
node_config = NodeConfig('uh')
|
|
node_config.host = "111.222.111.222"
|
|
node_config.port = 1122
|
|
pgbouncer_config = {"host": "192.168.0.1", "port": 7654}
|
|
pg = PgConnectionViaPgbouncer(
|
|
node_config, pgbouncer_config, Logger(), stub_psycopg2)
|
|
|
|
ok, err, cursor = pg.verify_connection()
|
|
self.assertIn("inet_server_addr()", cursor.query)
|
|
self.assertIn("%(host)s", cursor.query)
|
|
self.assertIn("inet_server_port()", cursor.query)
|
|
self.assertIn("%(port)s", cursor.query)
|
|
self.assertEqual({
|
|
"host": "111.222.111.222",
|
|
"port": 1122
|
|
}, cursor.params)
|
|
|
|
def test_WhenVerifyQueryReturnsIssue_IssueIsReported(self):
|
|
cursor = StubCursor(('SELECT', [['Bokito has escaped']]))
|
|
conn = StubConnection(cursor)
|
|
stub_psycopg2 = StubPsycopg2(conn)
|
|
node_config = NodeConfig('xyz123')
|
|
node_config.host = "111.222.111.222"
|
|
node_config.port = 1122
|
|
pgbouncer_config = {"host": "192.168.0.1", "port": 7654}
|
|
pg = PgConnectionViaPgbouncer(
|
|
node_config, pgbouncer_config, Logger(), stub_psycopg2)
|
|
|
|
ok, err, _ = pg.verify_connection()
|
|
self.assertFalse(ok)
|
|
self.assertIn("not connected to the expected PostgreSQL backend", err)
|
|
self.assertIn("Bokito has escaped", err)
|
|
|
|
def test_WhenVerifyQueryReturnsNoIssue_OkIsReported(self):
|
|
conn = StubConnection(StubCursor(('SELECT', [[None]])))
|
|
stub_psycopg2 = StubPsycopg2(conn)
|
|
node_config = NodeConfig('xyz123')
|
|
pgbouncer_config = {"host": "192.168.0.1", "port": 7654}
|
|
pg = PgConnectionViaPgbouncer(
|
|
node_config, pgbouncer_config, Logger(), stub_psycopg2)
|
|
|
|
ok, err, _ = pg.verify_connection()
|
|
self.assertTrue(ok)
|
|
self.assertIs(None, err)
|
|
|
|
class PgBouncerConsoleConnectionTests(unittest.TestCase):
|
|
def test_OnConnect_AutoCommitIsEnabled(self):
|
|
conn = StubConnection()
|
|
stub_psycopg2 = StubPsycopg2().add_connection(conn)
|
|
pg = PgBouncerConsoleConnection(
|
|
NodeConfig("bob"), Logger(), stub_psycopg2)
|
|
|
|
self.assertFalse(conn.autocommit)
|
|
pg.connect()
|
|
self.assertTrue(conn.autocommit)
|