diff --git a/src/rest_search/__init__.py b/src/rest_search/__init__.py index 9e21db9..034416a 100644 --- a/src/rest_search/__init__.py +++ b/src/rest_search/__init__.py @@ -3,7 +3,7 @@ from threading import local from aws_requests_auth.aws_auth import AWSRequestsAuth -from botocore.session import Session +from aws_requests_auth.boto_utils import BotoAWSRequestsAuth from django.conf import settings from opensearchpy import OpenSearch, RequestsHttpConnection @@ -58,14 +58,8 @@ def __create_connection(self, config): aws_service="es", ) elif "AWS_REGION" in config: - session = Session() - creds = session.get_credentials().get_frozen_credentials() - kwargs["connection_class"] = RequestsHttpConnection - kwargs["http_auth"] = AWSRequestsAuth( - aws_access_key=creds.access_key, - aws_secret_access_key=creds.secret_key, - aws_token=creds.token, + kwargs["http_auth"] = BotoAWSRequestsAuth( aws_host=config["HOST"], aws_region=config["AWS_REGION"], aws_service="es", diff --git a/tests/test_connection.py b/tests/test_connection.py index 5b8774c..56fcb26 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -from unittest.mock import Mock, patch +from unittest.mock import patch from django.test import TestCase, override_settings @@ -42,15 +42,7 @@ def test_aws_auth(self, mock_opensearch): } ) @patch("rest_search.OpenSearch") - @patch("rest_search.Session") - def test_aws_role_auth(self, mock_session, mock_opensearch): - mock_creds = Mock( - access_key="mock-access-key", - secret_key="mock-secret-key", - token="mock-token", - ) - mock_session.return_value.get_credentials.return_value.get_frozen_credentials.return_value = mock_creds - + def test_aws_role_auth(self, mock_opensearch): es = connections["default"] self.assertIsNotNone(es)