def __init__(self, cur: Optional[MockCursor] = None, connection: Optional[MockPsycopgConnection] = None, version: str = '90602', name: str = 'postgres', host: str = 'localhost', port: str = '25565', user: str = 'postgres'): # Setup mocks for the connection self.close = mock.MagicMock() self.cursor = mock.MagicMock(return_value=cur) # if no mock pyscopg connection passed, create default one if not connection: connection = MockPsycopgConnection(cursor=cur, dsn_parameters={ 'dbname': name, 'host': host, 'port': port, 'user': user }) # mock psycopg2.connect call in PostgreSQLConnection.__init__ to return mock psycopg connection with mock.patch('psycopg2.connect', mock.Mock(return_value=connection)): super().__init__({ "host_name": host, "user_name": user, "port": port, "database_name": name })
def test_execute_dict_success(self): # Setup: Create a mock server connection that will return a result set mock_cursor = utils.MockCursor(utils.get_mock_results()) mock_conn = MockPsycopgConnection(cursor=mock_cursor) # noinspection PyTypeChecker with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_conn)): server_conn = PostgreSQLConnection({}) # If: I execute a query as a dictionary results = server_conn.execute_dict('SELECT * FROM pg_class') # Then: # ... Both the columns and the results should be returned self.assertIsInstance(results, tuple) self.assertEqual(len(results), 2) # ... I should get a list of columns returned to me cols = results[0] self.assertIsInstance(cols, list) self.assertListEqual(cols, mock_cursor.description) # ... I should get the results formatted as a list of dictionaries rows = results[1] self.assertIsInstance(rows, list) for idx, row in enumerate(rows): self.assertDictEqual(row, mock_cursor._results[1][idx]) # ... The cursor should be closed mock_cursor.close.assert_called_once()
def setUp(self): """Set up the tests with common connection parameters""" # Set up the mock connection service and connection info self.connection_service = ConnectionService() self.connection_service._service_provider = utils.get_mock_service_provider( {WORKSPACE_SERVICE_NAME: WorkspaceService()}) self.owner_uri = 'test_uri' self.connection_type = ConnectionType.DEFAULT self.connect_params: ConnectRequestParams = ConnectRequestParams.from_dict( { 'ownerUri': self.owner_uri, 'type': self.connection_type, 'connection': { 'options': {} } }) self.mock_psycopg_connection = MockPsycopgConnection(dsn_parameters={ 'host': 'myserver', 'dbname': 'postgres', 'user': '******' }) # Mock psycopg2's connect method to store the current cancellation token. This lets us # capture the cancellation token state as it would be during a long-running connection. self.token_store = []
def test_server_conn_init(self): # Setup: Create a mock connection with an 'interesting' version dsn_parameters = {'dbname': 'postgres', 'host': 'localhost', 'port': '25565', 'user': '******'} mock_conn = MockPsycopgConnection(dsn_parameters=dsn_parameters) mock_conn.server_version = '100216' # If: I initialize a server connection # noinspection PyTypeChecker with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_conn)): server_conn = PostgreSQLConnection({}) # Then: The properties should be properly set self.assertEqual(server_conn._conn, mock_conn) self.assertEqual(server_conn.connection, mock_conn) self.assertDictEqual(server_conn._dsn_parameters, dsn_parameters) self.assertTupleEqual((10, 2, 16), server_conn.server_version)
def server_info_is_cloud_internal(self, host_suffix, is_cloud): """Test that the connection response handles cloud connections correctly""" # Set up the parameters for the connection connection_uri = 'someuri' connection_details = ConnectionDetails() connection_details.options = { 'user': '******', 'password': '******', 'host': f'myserver{host_suffix}', 'dbname': 'postgres' } connection_type = ConnectionType.DEFAULT # Set up the mock connection for psycopg2's connect method to return mock_connection = MockPsycopgConnection( dsn_parameters={ 'host': f'myserver{host_suffix}', 'dbname': 'postgres', 'user': '******' }) # Set up the connection service and call its connect method with the # supported options with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_connection)): response = self.connection_service.connect( ConnectRequestParams(connection_details, connection_uri, connection_type)) # Verify that the response's serverInfo.isCloud attribute is set correctly self.assertIsNotNone(response.connection_id) self.assertIsNotNone(response.server_info.server_version) self.assertEqual(response.server_info.is_cloud, is_cloud)
def test_get_connection_for_existing_connection(self): """Test that get_connection returns a connection that already exists for the given URI and type""" # Set up the test with mock data connection_uri = 'someuri' connection_type = ConnectionType.EDIT mock_connection = MockPsycopgConnection(dsn_parameters={ 'host': 'myserver', 'dbname': 'postgres', 'user': '******' }) # Insert a ConnectionInfo object into the connection service's map connection_details = ConnectionDetails.from_data({}) connection_info = ConnectionInfo(connection_uri, connection_details) self.connection_service.owner_to_connection_map[ connection_uri] = connection_info # Get the connection without first creating it with mock.patch( 'psycopg2.connect', new=mock.Mock( return_value=mock_connection)) as mock_psycopg2_connect: connection = self.connection_service.get_connection( connection_uri, connection_type) mock_psycopg2_connect.assert_called_once() self.assertEqual(connection._conn, mock_connection)
def test_non_default_database(self): """Test that if a database is given, the default database is not used""" # Set up the connection params and default database name default_db = 'test_db' actual_db = 'postgres' self.connection_service._service_provider[ WORKSPACE_SERVICE_NAME].configuration.pgsql.default_database = default_db params: ConnectRequestParams = ConnectRequestParams.from_dict({ 'ownerUri': 'someUri', 'type': ConnectionType.DEFAULT, 'connection': { 'options': { 'user': '******', 'password': '******', 'host': 'myserver', 'dbname': actual_db } } }) # If I connect with an empty database name with mock.patch('ossdbtoolsservice.connection.connection_service._build_connection_response'), \ mock.patch('psycopg2.connect', return_value=MockPsycopgConnection()) as mock_psycopg2_connect: self.connection_service.connect(params) # Then psycopg2's connect method was called with the default database calls = mock_psycopg2_connect.mock_calls self.assertEqual(len(calls), 1) self.assertNotEqual(calls[0][2]['dbname'], default_db) self.assertEqual(calls[0][2]['dbname'], actual_db)
def test_list_databases(self): """Test that the list databases handler correctly lists the connection's databases""" # Set up the test with mock data mock_query_results = [('database1', ), ('database2', )] connection_uri = 'someuri' mock_psycopg_connection = MockPsycopgConnection( dsn_parameters={ 'host': 'myserver', 'dbname': 'postgres', 'user': '******' }, cursor=MockCursor(mock_query_results)) mock_request_context = utils.MockRequestContext() # Insert a ConnectionInfo object into the connection service's map connection_details = ConnectionDetails.from_data({}) connection_info = ConnectionInfo(connection_uri, connection_details) self.connection_service.owner_to_connection_map[ connection_uri] = connection_info # Verify that calling the listdatabases handler returns the expected databases params = ListDatabasesParams() params.owner_uri = connection_uri with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_psycopg_connection)): self.connection_service.handle_list_databases( mock_request_context, params) expected_databases = [result[0] for result in mock_query_results] self.assertEqual( mock_request_context.last_response_params.database_names, expected_databases)
def setUp(self): self._metadata_factory = mock.MagicMock() self._mock_cursor = MockCursor(None) self._connection = MockPsycopgConnection({"port": "8080", "host": "test", "dbname": "test"}, self._mock_cursor) self._initialize_edit_request = InitializeEditParams() self._initialize_edit_request.schema_name = 'public' self._initialize_edit_request.object_name = 'Employee' self._initialize_edit_request.object_type = 'Table' db_column = DbColumn() column = EditColumnMetadata(db_column, None) self._columns_metadata = [column] self._schema_name = 'public' self._table_name = 'table' self._edit_table_metadata = EditTableMetadata(self._schema_name, self._table_name, self._columns_metadata) self._query_executer = mock.MagicMock() self._on_success = mock.MagicMock() self._on_failure = mock.MagicMock() self._data_editor_session = DataEditorSession(self._metadata_factory) self._metadata_factory.get = mock.Mock(return_value=self._edit_table_metadata) self._query = 'SELECT TESTCOLUMN FROM TESTTABLE LIMIT 100' self._data_editor_session._construct_initialize_query = mock.Mock(return_value=self._query)
def test_execute_dict_fail(self): # Setup: Create a mock psycopg connection that will raise an exception mock_cursor = utils.MockCursor(None, throw_on_execute=True) mock_conn = MockPsycopgConnection(cursor=mock_cursor) # noinspection PyTypeChecker with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_conn)): server_conn = PostgreSQLConnection({}) # If: I execute a query as a dictionary # Then: # ... I should get an exception with self.assertRaises(Exception): server_conn.execute_dict('SELECT * FROM pg_class') # ... The cursor should be closed mock_cursor.close.assert_called_once()
def test_same_options_uses_existing_connection(self): """Test that the connect method uses an existing connection when connecting again with the same options""" # Set up the test with mock data connection_uri = 'someuri' connection_type = ConnectionType.DEFAULT mock_psycopg_connection = MockPsycopgConnection(dsn_parameters={ 'host': 'myserver', 'dbname': 'postgres', 'user': '******' }) mock_server_connection = MockPGServerConnection(cur=None, host='myserver', name='postgres', user='******') # Insert a ConnectionInfo object into the connection service's map old_connection_details = ConnectionDetails.from_data({ 'host': 'myserver', 'dbname': 'postgres', 'user': '******', 'abc': 123 }) old_connection_info = ConnectionInfo(connection_uri, old_connection_details) old_connection_info.add_connection(connection_type, mock_server_connection) self.connection_service.owner_to_connection_map[ connection_uri] = old_connection_info # Connect with identical options, and verify that disconnect was not called params: ConnectRequestParams = ConnectRequestParams.from_dict({ 'ownerUri': connection_uri, 'type': connection_type, 'connection': { 'options': old_connection_details.options } }) with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_psycopg_connection) ) as mock_psycopg2_connect: response = self.connection_service.connect(params) mock_psycopg2_connect.assert_not_called() mock_psycopg_connection.close.assert_not_called() self.assertIsNotNone(response.connection_id)
def test_recovery_properties(self): # Setup: # NOTE: We're *not* mocking out the template rendering b/c this will verify that there's a template # ... Create a mock query execution that will return the properties mock_exec_dict = mock.MagicMock(return_value=([], [TestServer.CHECK_RECOVERY_ROW])) # ... Create an instance of the class and override the connection mock_connection = MockPsycopgConnection({'host': 'host', 'dbname': 'dbname'}) with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_connection)): pg_connection = PostgreSQLConnection({}) pg_connection.execute_dict = mock_exec_dict obj = Server(pg_connection) # If: I retrieve all the values in the recovery properties # Then: # ... The properties based on the properties should be availble self.assertEqual(obj.in_recovery, TestServer.CHECK_RECOVERY_ROW['inrecovery']) self.assertEqual(obj.wal_paused, TestServer.CHECK_RECOVERY_ROW['isreplaypaused'])
def run_on_connect_callback(self, conn_type: ConnectionType, expect_callback: bool) -> None: """Inner function for callback tests that verifies expected behavior given different connection types""" callbacks = [MagicMock(), MagicMock()] for callback in callbacks: self.connection_service.register_on_connect_callback(callback) # Set up the parameters for the connection connection_uri = 'someuri' connection_details = ConnectionDetails() connection_details.options = { 'user': '******', 'password': '******', 'host': f'myserver', 'dbname': 'postgres' } connection_type = conn_type # Set up the mock connection for psycopg2's connect method to return mock_psycopg_connection = MockPsycopgConnection( dsn_parameters={ 'host': f'myserver', 'dbname': 'postgres', 'user': '******' }) # Set up the connection service and call its connect method with the # supported options with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_psycopg_connection)): self.connection_service.connect( ConnectRequestParams(connection_details, connection_uri, connection_type)) self.connection_service.get_connection(connection_uri, conn_type) # ... The mock config change callbacks should have been called for callback in callbacks: if (expect_callback): callback.assert_called_once() # Verify call args match expected callargs: ConnectionInfo = callback.call_args[0][0] self.assertEqual(callargs.owner_uri, connection_uri) else: callback.assert_not_called()
def test_connect(self): """Test that the service connects to a PostgreSQL server""" # Set up the parameters for the connection params: ConnectRequestParams = ConnectRequestParams.from_dict({ 'ownerUri': 'someUri', 'type': ConnectionType.DEFAULT, 'connection': { 'options': { 'user': '******', 'password': '******', 'host': 'myserver', 'dbname': 'postgres' } } }) # Set up the mock connection for psycopg2's connect method to return mock_connection = MockPsycopgConnection(dsn_parameters={ 'host': 'myserver', 'dbname': 'postgres', 'user': '******' }) # Set up the connection service and call its connect method with the supported options with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_connection)): response = self.connection_service.connect(params) # Verify that psycopg2's connection method was called and that the # response has a connection id, indicating success. self.assertIs( self.connection_service.owner_to_connection_map[ params.owner_uri].get_connection(params.type)._conn, mock_connection) self.assertIsNotNone(response.connection_id) self.assertIsNotNone(response.server_info.server_version) self.assertFalse(response.server_info.is_cloud)