From 553d0075f669e5b3dcba08fedcea0466d91bfc35 Mon Sep 17 00:00:00 2001 From: Divya Madala Date: Mon, 18 Mar 2024 15:01:26 -0700 Subject: [PATCH] Final changes Signed-off-by: Divya Madala --- .../benchmark_test/benchmark_create_cluster.py | 4 ++-- .../benchmark_test/benchmark_test_cluster.py | 3 ++- .../benchmark_test_runner_opensearch.py | 2 +- .../benchmark_test/test_benchmark_test_cluster.py | 1 - .../test_benchmark_test_runner_opensearch.py | 11 ++++------- 5 files changed, 9 insertions(+), 12 deletions(-) diff --git a/src/test_workflow/benchmark_test/benchmark_create_cluster.py b/src/test_workflow/benchmark_test/benchmark_create_cluster.py index b36b974991..946ce47228 100644 --- a/src/test_workflow/benchmark_test/benchmark_create_cluster.py +++ b/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 diff --git a/src/test_workflow/benchmark_test/benchmark_test_cluster.py b/src/test_workflow/benchmark_test/benchmark_test_cluster.py index 42f71a9779..42e29f39c0 100644 --- a/src/test_workflow/benchmark_test/benchmark_test_cluster.py +++ b/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 diff --git a/src/test_workflow/benchmark_test/benchmark_test_runner_opensearch.py b/src/test_workflow/benchmark_test/benchmark_test_runner_opensearch.py index 1c8cab1de9..56d0b5425c 100644 --- a/src/test_workflow/benchmark_test/benchmark_test_runner_opensearch.py +++ b/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 diff --git a/tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_cluster.py b/tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_cluster.py index 81295d0977..ae5163e5ff 100644 --- a/tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_cluster.py +++ b/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" diff --git a/tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_runner_opensearch.py b/tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_runner_opensearch.py index 524650543e..62736438db 100644 --- a/tests/tests_test_workflow/test_benchmark_workflow/benchmark_test/test_benchmark_test_runner_opensearch.py +++ b/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) -