taskbuffer/tests/db_connection_tests.py

0001 # Description: Unit tests for the call timeout of DBProxy connections. No database is used:
0002 # the drivers are replaced
0003 import unittest
0004 from typing import Any, cast
0005 from unittest import mock
0006
0007 from pandaserver.config import panda_config
0008 from pandaserver.taskbuffer import OraDBProxy, WrappedPostgresConn, wrapped_oracle_conn
0009
0010
0011 class FakeDriverConnection:
0012     # a plain object, so that an attribute the code under test never set is absent. Declared
0013     # only, the way the driver's connection has it
0014     call_timeout: int
0015
0016     def __init__(self, **kwargs: Any) -> None:
0017         self.kwargs = kwargs
0018         self.closed = False
0019
0020     def close(self) -> None:
0021         self.closed = True
0022
0023
0024 class FakeCursor:
0025     def __init__(self, conn: Any) -> None:
0026         self.conn = conn
0027         self.executed: list[str] = []
0028
0029     def initialize(self) -> str:
0030         return "db-host"
0031
0032     def execute(self, sql: str, var_map: dict[str, Any] | None = None) -> None:
0033         self.executed.append(sql)
0034
0035
0036 class ProxyTestBase(unittest.TestCase):
0037     backend = "oracle"
0038
0039     def setUp(self) -> None:
0040         self.driver_connections: list[FakeDriverConnection] = []
0041
0042         def driver_connect(**kwargs: Any) -> FakeDriverConnection:
0043             conn = FakeDriverConnection(**kwargs)
0044             self.driver_connections.append(conn)
0045             return conn
0046
0047         fake_oracledb = mock.Mock(connect=mock.Mock(side_effect=driver_connect))
0048         fake_psycopg = mock.Mock(connect=mock.Mock(side_effect=driver_connect))
0049         # create=True since the module imports only the driver of the configured backend
0050         patches: list[Any] = [
0051             mock.patch.object(OraDBProxy, "oracledb", fake_oracledb, create=True),
0052             mock.patch.object(OraDBProxy, "psycopg", fake_psycopg, create=True),
0053             mock.patch.object(OraDBProxy, "wrapped_oracle_conn", wrapped_oracle_conn, create=True),
0054             mock.patch.object(OraDBProxy, "WrappedPostgresConn", WrappedPostgresConn, create=True),
0055             mock.patch.object(OraDBProxy, "WrappedCursor", FakeCursor),
0056             mock.patch.object(OraDBProxy.atexit, "register"),
0057             mock.patch.object(panda_config, "dump_sql", False),
0058         ]
0059         for patch in patches:
0060             patch.start()
0061             self.addCleanup(patch.stop)
0062         self.proxy = OraDBProxy.DBProxy()
0063         self.proxy.backend = self.backend
0064
0065     def connect(self, call_timeout_ms: int | None) -> bool:
0066         return self.proxy.connect(dbhost="host", dbpasswd="secret", dbuser="reader", dbname="db", dbtimeout=60, dbport=5432, call_timeout_ms=call_timeout_ms)
0067
0068
0069 class TestOracleCallTimeout(ProxyTestBase):
0070     backend = "oracle"
0071
0072     def test_set_on_driver_connection_and_kept_on_reconnect(self) -> None:
0073         self.assertTrue(self.connect(30000))
0074         self.assertEqual(self.proxy.conn.orig_conn.call_timeout, 30000)
0075         self.assertTrue(self.proxy.connect(reconnect=True))
0076         self.assertEqual(len(self.driver_connections), 2)
0077         self.assertTrue(self.driver_connections[0].closed)
0078         self.assertIs(self.proxy.conn.orig_conn, self.driver_connections[1])
0079         self.assertEqual(self.driver_connections[1].call_timeout, 30000)
0080         self.assertEqual(self.driver_connections[1].kwargs["user"], "reader")
0081
0082     def test_not_set_without_limit(self) -> None:
0083         self.assertTrue(self.connect(None))
0084         self.assertFalse(hasattr(self.driver_connections[0], "call_timeout"))
0085
0086     def test_lowered_within_limit(self) -> None:
0087         self.assertTrue(self.connect(30000))
0088         raw = self.proxy.conn.orig_conn
0089         self.proxy.set_call_timeout(5000)
0090         self.assertEqual(raw.call_timeout, 5000)
0091         self.proxy.set_call_timeout(60000)
0092         self.assertEqual(raw.call_timeout, 30000)
0093         self.proxy.set_call_timeout(None)
0094         self.assertEqual(raw.call_timeout, 30000)
0095         # 0 would mean no limit to the driver
0096         with self.assertRaises(ValueError):
0097             self.proxy.set_call_timeout(0)
0098         self.assertEqual(raw.call_timeout, 30000)
0099
0100     def test_put_back_to_no_limit(self) -> None:
0101         self.assertTrue(self.connect(None))
0102         self.proxy.set_call_timeout(5000)
0103         self.assertEqual(self.proxy.conn.orig_conn.call_timeout, 5000)
0104         self.proxy.set_call_timeout(None)
0105         self.assertEqual(self.proxy.conn.orig_conn.call_timeout, 0)
0106
0107
0108 class TestPostgresCallTimeout(ProxyTestBase):
0109     backend = "postgres"
0110
0111     def test_in_dsn_and_kept_on_reconnect(self) -> None:
0112         self.assertTrue(self.connect(30000))
0113         self.assertTrue(self.proxy.connect(reconnect=True))
0114         self.assertEqual(len(self.driver_connections), 2)
0115         for conn in self.driver_connections:
0116             self.assertEqual(conn.kwargs["options"], "-c statement_timeout=30000")
0117             self.assertEqual(conn.kwargs["user"], "reader")
0118
0119     def test_not_in_dsn_without_limit(self) -> None:
0120         self.assertTrue(self.connect(None))
0121         self.assertNotIn("options", self.driver_connections[0].kwargs)
0122
0123     def test_lowered_with_set_local(self) -> None:
0124         self.assertTrue(self.connect(30000))
0125         self.proxy.set_call_timeout(5000)
0126         self.proxy.set_call_timeout(60000)
0127         self.proxy.set_call_timeout(None)
0128         executed = [sql.split(" /*")[0] for sql in cast(FakeCursor, self.proxy.cur).executed]
0129         # the end of the transaction puts the limit back, so None runs nothing
0130         self.assertEqual(executed, ["SET LOCAL statement_timeout TO 5000", "SET LOCAL statement_timeout TO 30000"])
0131
0132
0133 class TestMySQLCallTimeout(ProxyTestBase):
0134     backend = "mysql"
0135
0136     def test_refused(self) -> None:
0137         self.assertFalse(self.connect(30000))
0138         self.assertEqual(self.driver_connections, [])
0139
0140
0141 if __name__ == "__main__":
0142     unittest.main()