Final changes

Signed-off-by: Divya Madala <[email protected]>

src/test_workflow/benchmark_test/benchmark_create_cluster.py

@@ -15,7 +15,6 @@
 
 from manifests.build_manifest import BuildManifest
 from manifests.bundle_manifest import BundleManifest
-
 from test_workflow.benchmark_test.benchmark_args import BenchmarkArgs
 from test_workflow.benchmark_test.benchmark_test_cluster import BenchmarkTestCluster
 from test_workflow.integ_test.utils import get_password
@@ -61,7 +60,8 @@ def __init__(
                 else:
                     params_list.append(f" -c {key}={value}")
         role_params = (
-            " --require-approval=never"
+            f" --require-approval=never --plugin cdk-assume-role-credential-plugin"
+            f" -c assume-role-credentials:writeIamRoleName={role} -c assume-role-credentials:readIamRoleName={role} "
         )
         self.params = "".join(params_list) + role_params
         self.is_endpoint_public = False

src/test_workflow/benchmark_test/benchmark_test_cluster.py

@@ -6,10 +6,11 @@
 # compatible open source license.
 
 
+import json
 import logging
 import subprocess
+
 import requests
-import json
 from requests.auth import HTTPBasicAuth
 from retry.api import retry_call  # type: ignore
 

src/test_workflow/benchmark_test/benchmark_test_runner_opensearch.py

@@ -18,8 +18,8 @@
 from system.temporary_directory import TemporaryDirectory
 from system.working_directory import WorkingDirectory
 from test_workflow.benchmark_test.benchmark_args import BenchmarkArgs
-from test_workflow.benchmark_test.benchmark_test_cluster import BenchmarkTestCluster
 from test_workflow.benchmark_test.benchmark_create_cluster import BenchmarkCreateCluster
+from test_workflow.benchmark_test.benchmark_test_cluster import BenchmarkTestCluster
 from test_workflow.benchmark_test.benchmark_test_runner import BenchmarkTestRunner
 from test_workflow.benchmark_test.benchmark_test_suite import BenchmarkTestSuite
 

tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_cluster.py

@@ -73,7 +73,6 @@ def test_endpoint_with_security(self, mock_http_auth: Mock, mock_requests_get: M
         self.assertEqual(self.benchmark_test_cluster.endpoint_with_port, 'opensearch-cluster.amazon.com:443')
         self.assertEqual(self.benchmark_test_cluster.port, 443)
 
-
     def test_endpoint_with_timeout_error(self) -> None:
         self.args.insecure = True
         self.args.cluster_endpoint = "opensearch-cluster.amazon.com"

tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_runner_opensearch.py

@@ -7,9 +7,8 @@
 import os
 import tempfile
 import unittest
-from typing import Optional
-from typing import Any
-from unittest.mock import Mock, patch, MagicMock
+from typing import Any, Optional
+from unittest.mock import MagicMock, Mock, patch
 
 from manifests.bundle_manifest import BundleManifest
 from test_workflow.benchmark_test.benchmark_args import BenchmarkArgs
@@ -81,8 +80,6 @@ def test_run_with_dist_url_and_version(self, mock_suite: Mock, mock_cluster: Moc
     @patch('test_workflow.benchmark_test.benchmark_test_runner_opensearch.retry_call')
     def test_run_with_cluster_endpoint(self, mock_retry_call: Mock, mock_suite: Mock, mock_benchmark_test_cluster: Mock) -> None:
         args = MagicMock(cluster_endpoint=True)
-
-
         mock_cluster = MagicMock()
 
         mock_benchmark_test_cluster.return_value = mock_cluster
@@ -97,7 +94,8 @@ def test_run_with_cluster_endpoint(self, mock_retry_call: Mock, mock_suite: Mock
     @patch('test_workflow.benchmark_test.benchmark_test_runner_opensearch.retry_call')
     @patch("subprocess.run")
     @patch("requests.get")
-    def test_run_with_cluster_endpoint_with_arguments(self, mock_requests_get: Mock, mock_subprocess_run: Mock, mock_retry_call: Mock, mock_suite: Mock, mock_wait_for_processing: Optional[Mock]) -> None:
+    def test_run_with_cluster_endpoint_with_arguments(self, mock_requests_get: Mock, mock_subprocess_run: Mock,
+                                                      mock_retry_call: Mock, mock_suite: Mock, mock_wait_for_processing: Optional[Mock]) -> None:
         args = MagicMock(cluster_endpoint=True)
         mock_wait_for_processing.return_value = None
         mock_result = MagicMock()
@@ -126,4 +124,3 @@ def test_run_with_cluster_endpoint_with_arguments(self, mock_requests_get: Mock,
         self.assertEqual(mock_suite.call_count, 1)
         self.assertEqual(MockBenchmarkTestCluster.call_count, 1)
         mock_retry_call.assert_called_once_with(mock_suite.return_value.execute, tries=3, delay=60, backoff=2)
-