"""Test Opensearch connector.""" from unittest.mock import call, NonCallableMagicMock from unittest.mock import MagicMock from contextlib import nullcontext as does_not_raise from unittest.mock import patch import opensearchpy.exceptions import pytest from content_utils.connectors import opensearch as os_connector from content_utils.exceptions import IndexingFailedError from content_utils.exceptions import IneligibleEventError @patch('content_utils.connectors.opensearch.OpenSearch') def test_index_product_success(mock_os, mock_index_document): """Test adding product to index successfully.""" mock_os.return_value.index.return_value = { 'result': 'created', 'created': mock_index_document } test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), username='foo', password='oa:1') test_conn.index_product(mock_index_document) mock_os.return_value.index.assert_called_once_with( index='review_write', id=1, body=mock_index_document ) @patch('content_utils.connectors.opensearch.OpenSearch') def test_d3_subaccount_index_product_success( mock_os, mock_d3_subaccount_index_document): """Test adding d3 subaccount product to index successfully.""" mock_os.return_value.index.return_value = { 'result': 'created', 'created': mock_d3_subaccount_index_document } test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), username='foo', password='oa:1') test_conn.index_product(mock_d3_subaccount_index_document) mock_os.return_value.index.assert_called_once_with( index='review_write', id=1, body=mock_d3_subaccount_index_document ) @patch('content_utils.connectors.opensearch.OpenSearch') def test_index_product_fail(mock_os, mock_index_document): """Test adding product to index fail raises proper error.""" mock_os.return_value.index.return_value = {'result': 'not created'} test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), username='foo', password='oa:1') with pytest.raises(IndexingFailedError) as err: test_conn.index_product(mock_index_document) mock_os.return_value.index.assert_called_once_with( index='review_write', id=1, body=mock_index_document ) assert str(err.value) == 'Failed to index product.' @pytest.mark.parametrize( ( 'test_description', 'opensearch_response', 'expected_raise', 'expected_raise_message', ), [ ( 'success', { '_index': 'review_write', '_type': '_doc', '_id': '12321', '_version': 5, 'result': 'updated', '_shards': { 'total': 2, 'successful': 2, 'failed': 0 }, '_seq_no': 14, '_primary_term': 2 }, does_not_raise(), 'None', ), ( 'noop, not found', { 'status': 404 }, does_not_raise(), 'None', ), ( 'noop, no updates', { '_index': 'review_write', '_type': '_doc', '_id': '12321', '_version': 4, 'result': 'noop', '_shards': { 'total': 0, 'successful': 0, 'failed': 0 }, '_seq_no': 71, '_primary_term': 1 }, does_not_raise(), 'None', ), ( 'failure', { '_index': 'review_write', '_type': '_doc', '_id': '12321', '_version': 5, 'result': 'not updated', '_shards': { 'total': 2, 'successful': 0, 'failed': 2 }, '_seq_no': 14, '_primary_term': 2 }, pytest.raises(IndexingFailedError), 'Failed to patch product. ID 12321', ), ], ) def test_patch_product( mocker, test_description, opensearch_response, expected_raise, expected_raise_message, ): """Test patch_product.""" mocker.patch.object( os_connector, 'OpenSearch', return_value=NonCallableMagicMock( update=MagicMock( return_value=opensearch_response ) ) ) test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), username='foo', password='oa:1') with expected_raise as er: test_conn.patch_product({ 'product_id': 12321, 'locked_until_datetime': '2022-07-18T19:42:22.872582Z' }) assert str(getattr(er, 'value', None)) == expected_raise_message assert os_connector.OpenSearch.return_value.update.mock_calls == [ call( index='review_write', id=12321, ignore=[404], body={ 'doc': { 'product_id': 12321, 'locked_until_datetime': '2022-07-18T19:42:22.872582Z' } } ) ] @patch('content_utils.connectors.opensearch.OpenSearch') def test_drop_index(mock_os): """Test drop index.""" mock_os.return_value.indices_exists_alias.return_value = True mock_os.return_value.aliases.delete.return_value = {'result': 'not created'} test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), username='foo', password='oa:1') test_conn.drop_index() mock_os.return_value.aliases.delete.assert_called_once_with( name='review_write', ignore=[400, 404] ) mock_os.return_value.indices.delete.assert_called_once_with( index='review.v01', ignore=[400, 404] ) @patch('content_utils.connectors.opensearch.OpenSearch') def test_clear_index_via_alias(mock_os): """Test clear_index_via_alias.""" mock_os.return_value.delete_by_query.return_value = None test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), username='foo', password='oa:1') test_conn.clear_index_via_alias() mock_os.return_value.delete_by_query.assert_called_once_with( index='review_write', body={ 'query': { 'match_all': {} } }, refresh=True, wait_for_completion=True, ignore=[400, 404] ) @patch('content_utils.connectors.opensearch.OpenSearch') def test_remove_product(mock_os): """Test remove_product.""" mock_os.return_value.delete.return_value = None test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), username='foo', password='oa:1') test_conn.remove_product('123') mock_os.return_value.delete.assert_called_once_with( index='review_write', id='123', ignore=[404] ) @patch('content_utils.connectors.opensearch.OpenSearch') def test_check_product_existing_status_found_success( mock_os, mock_row_data, mock_index_document): """Test check_product_existing_status.""" mock_os.return_value.get.return_value = { '_source': mock_index_document, 'found': True } test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), username='foo', password='oa:1') response = test_conn.check_product_existing_status(mock_row_data, expected_status='found') assert response == mock_index_document @patch('content_utils.connectors.opensearch.OpenSearch') def test_check_product_existing_status_found_failure(mock_os, mock_row_data): """Test check_product_existing_status.""" mock_err = MagicMock(side_effect=opensearchpy.exceptions.NotFoundError) mock_os.return_value.get.side_effect = mock_err test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), username='foo', password='oa:1') with pytest.raises(IneligibleEventError) as err: test_conn.check_product_existing_status(mock_row_data, expected_status='found') assert str(err.value) == 'Product not found in queue: product id: 1' @patch('content_utils.connectors.opensearch.OpenSearch') def test_check_product_existing_status_not_found_success(mock_os, mock_row_data): """Test check_product_existing_status.""" mock_err = MagicMock(side_effect=opensearchpy.exceptions.NotFoundError) mock_os.return_value.get.side_effect = mock_err test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), username='foo', password='oa:1') test_conn.check_product_existing_status(mock_row_data, expected_status='not found') @patch('content_utils.connectors.opensearch.OpenSearch') def test_check_product_existing_status_not_found_failure( mock_os, mock_row_data, mock_index_document): """Test check_product_existing_status.""" mock_os.return_value.get.return_value = { '_source': mock_index_document, 'found': True } test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), username='foo', password='oa:1') with pytest.raises(IneligibleEventError) as err: test_conn.check_product_existing_status(mock_row_data, expected_status='not found') assert str(err.value) == 'Product already in queue: queue id: 1 product id: 1' @patch('content_utils.connectors.opensearch.OpenSearch') @patch('content_utils.connectors.opensearch._get_assumed_role_credentials') def test_os_assume_role_auth( mock_creds, mock_os, mock_row_data, mock_index_document): """Test connecting via assumed role.""" mock_creds.return_value = { 'AccessKeyId': 'abc', 'SecretAccessKey': 'def', 'SessionToken': 'ghi', } mock_os.return_value.get.return_value = { '_source': mock_index_document, 'found': True } test_conn = os_connector.LambdaOpensearchConnector( 'url', MagicMock(), assume_role_arn='arn:aws:abc') response = test_conn.check_product_existing_status(mock_row_data, expected_status='found') assert response == mock_index_document @patch.dict('content_utils.connectors.opensearch.environ', {'AWS_REGION': 'foo'}) def test_get_current_region(): """Test get_current_region.""" actual = os_connector._get_current_region() assert actual == 'foo' @patch.dict( 'content_utils.connectors.opensearch.environ', {'AWS_DEFAULT_REGION': 'bar'}, clear=True ) def test_get_current_region_default(): """Test get_current_region.""" actual = os_connector._get_current_region() assert actual == 'bar' @patch('content_utils.connectors.opensearch.boto3') def test_get_assumed_role_credentials(mock_boto): """Test get_assumed_role_credentials.""" mock_boto.client.return_value.assume_role.return_value = { 'Credentials': 'foo' } actual = os_connector._get_assumed_role_credentials('foo') assert actual == 'foo' class MockCredentials: """Mock Creds.""" def __init__(self) -> None: """Init.""" self.access_key = 'foo' self.secret_key = 'bar' self.token = 'baz' @patch.dict('content_utils.connectors.opensearch.environ', {'AWS_REGION': 'foo'}) @patch('content_utils.connectors.opensearch.AWSRequestsAuth') @patch('content_utils.connectors.opensearch.OpenSearch') @patch('content_utils.connectors.opensearch.boto3') def test_os_connector_no_credentials(mock_boto, _, mock_auth): """Test creating the connector without creds.""" mock_boto.Session.return_value.get_credentials.return_value = MockCredentials() os_connector.LambdaOpensearchConnector('foo', MagicMock()) mock_auth.assert_called_once_with(**{ 'aws_access_key': 'foo', 'aws_host': 'foo', 'aws_region': 'foo', 'aws_secret_access_key': 'bar', 'aws_service': 'es', 'aws_token': 'baz' }) @patch('content_utils.connectors.opensearch.boto3') def test_bad_os_call(mock_boto): """Test calling an undefined method.""" mock_boto.client.return_value.assume_role.return_value = { 'Credentials': { 'AccessKeyId': 'foo', 'SecretAccessKey': 'bar', 'SessionToken': 'baz' } } mock_conn = os_connector.LambdaOpensearchConnector( 'foo', MagicMock(), assume_role_arn='arn:iam:us-east-1:123456:7890:role/foo') with pytest.raises(Exception) as ex: mock_conn._call('foo') assert ex.value == 'Unknown os method: foo' @patch('content_utils.connectors.opensearch.OpenSearch') @patch('content_utils.connectors.opensearch.boto3') def test_os_connector_retry(mock_boto, mock_os): """Test retrying when session is expired.""" mock_boto.client.return_value.assume_role.return_value = { 'Credentials': { 'AccessKeyId': 'foo', 'SecretAccessKey': 'bar', 'SessionToken': 'baz' } } mock_os.return_value.foo.side_effect = [ os_connector.opensearch_exceptions.AuthorizationException( 403, 'baz', {'message': 'expired'}), 'qux', os_connector.opensearch_exceptions.AuthorizationException( 403, 'baz', {'message': 'denied'})] mock_conn = os_connector.LambdaOpensearchConnector( 'foo', MagicMock(), assume_role_arn='arn:iam:us-east-1:123456:7890:role/foo') actual = mock_conn._call('foo') assert actual == 'qux' assert mock_boto.client.return_value.assume_role.call_count == 2 with pytest.raises(os_connector.opensearch_exceptions.AuthorizationException): mock_conn._call('foo')