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 = MockConnection(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), mock_connection) self.assertIsNotNone(response.connection_id) self.assertIsNotNone(response.server_info.server_version) self.assertFalse(response.server_info.is_cloud)
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[constants.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('pgsqltoolsservice.connection.connection_service._build_connection_response'), \ mock.patch('psycopg2.connect') 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_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_connection = MockConnection(dsn_parameters={ 'host': 'myserver', 'dbname': '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_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_connection)) as mock_psycopg2_connect: response = self.connection_service.connect(params) mock_psycopg2_connect.assert_not_called() mock_connection.close.assert_not_called() self.assertIsNotNone(response.connection_id)
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 = { constants.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_connection = MockConnection(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_changing_options_disconnects_existing_connection(self): """ Test that the connect method disconnects an existing connection when trying to open the same connection with different options """ # Set up the test with mock data connection_uri = 'someuri' connection_type = ConnectionType.DEFAULT mock_connection = MockConnection(dsn_parameters={ 'host': 'myserver', 'dbname': '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_connection) self.connection_service.owner_to_connection_map[ connection_uri] = old_connection_info # Create a different request with the same owner uri params: ConnectRequestParams = ConnectRequestParams.from_dict({ 'ownerUri': connection_uri, 'type': connection_type, 'connection': { 'options': { 'host': 'myserver', 'dbname': 'postgres', 'user': '******', 'abc': 234 } } }) # Connect with different options, and verify that disconnect was called with mock.patch('psycopg2.connect', new=mock.Mock(return_value=mock_connection)): self.connection_service.connect(params) mock_connection.close.assert_called_once()
def test_response_when_connect_fails(self): """Test that the proper response is given when a connection fails""" error_message = 'some error' params: ConnectRequestParams = ConnectRequestParams.from_dict({ 'ownerUri': 'someUri', 'type': ConnectionType.DEFAULT, 'connection': { 'options': { 'connectionString': '' } } }) with mock.patch('psycopg2.connect', new=mock.Mock(side_effect=Exception(error_message))): response = self.connection_service.connect(params) # The response should not have a connection ID and should contain the error message self.assertIsNone(response.connection_id) self.assertEqual(response.error_message, error_message)
def test_handle_connect_request(self): """Test that the handle_connect_request method kicks off a new thread to do the connection""" # Setup: Create a mock request context to handle output rc = utils.MockRequestContext() connect_response = ConnectionCompleteParams() self.connection_service.connect = Mock(return_value=connect_response) # If: I make a request to connect params: ConnectRequestParams = ConnectRequestParams.from_dict({ 'ownerUri': 'someUri', 'type': ConnectionType.QUERY, 'connection': { 'server_name': 'someserver', 'user_name': 'someuser', 'database_name': 'somedb', 'options': { 'password': '******' } } }) # Connect and wait for the thread to finish executing, then verify the connection information self.connection_service.handle_connect_request(rc, params) connection_thread = self.connection_service.owner_to_thread_map[ params.owner_uri] self.assertIsNotNone(connection_thread) connection_thread.join() # Then: # ... Connect should have been called once self.connection_service.connect.assert_called_once_with(params) # ... A True should have been sent as the response to the request rc.send_response.assert_called_once_with(True) # ... A connection complete notification should have been sent back as well rc.send_notification.assert_called_once_with( CONNECTION_COMPLETE_METHOD, connect_response) # ... An error should not have been called rc.send_error.assert_not_called()
def test_connect_with_access_token(self): """Test that the service connects to a PostgreSQL server using an access token as a password""" # Set up the parameters for the connection params: ConnectRequestParams = ConnectRequestParams.from_dict({ 'ownerUri': 'someUri', 'type': ConnectionType.DEFAULT, 'connection': { 'options': { 'user': '******', 'azureAccountToken': 'exampleToken', 'host': 'myserver', 'dbname': 'postgres' } } }) # Set up the mock connection for psycopg2's connect method to return mock_connection = MockConnection(dsn_parameters={ 'host': 'myserver', 'dbname': 'postgres', 'user': '******' }) # Set up psycopg2 instance for connection service to call mock_connect_method = mock.Mock(return_value=mock_connection) # Set up the connection service and call its connect method with the supported options with mock.patch('psycopg2.connect', new=mock_connect_method): response = self.connection_service.connect(params) # Verify that psycopg2's connection method was called with password set to account token. mock_connect_method.assert_called_once_with(user='******', password='******', host='myserver', dbname='postgres') # 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), mock_connection) self.assertIsNotNone(response.connection_id) self.assertIsNotNone(response.server_info.server_version) self.assertFalse(response.server_info.is_cloud)