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 = MockConnection(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 _connection_details() -> Tuple[ConnectionDetails, str]: param = ConnectionDetails() param.options = { 'host': TEST_HOST, 'dbname': TEST_DBNAME, 'user': TEST_USER } session_uri = ObjectExplorerService._generate_session_uri(param) return param, session_uri
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_connection = MockConnection(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_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()