Skip to content

Commit f412e7e

Browse files
committed
addressing more code review comments
1 parent c6d4409 commit f412e7e

File tree

1 file changed

+13
-39
lines changed

1 file changed

+13
-39
lines changed

tests/unit/test_iam_dsql_plugin.py

Lines changed: 13 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121

2222
import pytest
2323

24-
from aws_advanced_python_wrapper.errors import AwsWrapperError
2524
from aws_advanced_python_wrapper.hostinfo import HostInfo
2625
from aws_advanced_python_wrapper.iam_plugin import IamAuthPlugin, TokenInfo
2726
from aws_advanced_python_wrapper.utils.dsql_token_utils import DSQLTokenUtils
@@ -34,7 +33,7 @@
3433
_TEST_TOKEN = "test_token"
3534
_DEFAULT_PG_PORT = 5432
3635

37-
_PG_HOST_INFO = HostInfo("dsqltestclusternamefoobar1.dsql-gamma.us-east-2.on.aws")
36+
_PG_HOST_INFO = HostInfo("dsqltestclusternamefoobar1.dsql.us-east-2.on.aws")
3837
_PG_HOST_INFO_WITH_PORT = HostInfo(_PG_HOST_INFO.url, port=1234)
3938
_PG_REGION = "us-east-2"
4039

@@ -137,7 +136,7 @@ def test_pg_connect_valid_token_in_cache(user, mocker, mock_plugin_service, mock
137136
assert _GENERATED_TOKEN_NON_ADMIN != actual_token.token
138137

139138
assert _TEST_TOKEN == actual_token.token
140-
assert actual_token.is_expired() is False
139+
assert not actual_token.is_expired()
141140

142141

143142
@patch("aws_advanced_python_wrapper.iam_plugin.IamAuthPlugin._token_cache", _token_cache)
@@ -167,7 +166,7 @@ def test_pg_connect_with_invalid_port_fall_backs_to_host_port(
167166

168167
actual_token = _token_cache.get(f"{_PG_REGION}:{_PG_HOST_INFO.url}:1234:admin")
169168
assert _GENERATED_TOKEN == actual_token.token
170-
assert actual_token.is_expired() is False
169+
assert not actual_token.is_expired()
171170

172171
# Assert password has been updated to the value in token cache
173172
expected_props = {"user": "admin", "iam_default_port": "0"}
@@ -203,7 +202,7 @@ def test_pg_connect_with_invalid_port_and_no_host_port_fall_backs_to_host_port(
203202
actual_token = _token_cache.get(
204203
f"{_PG_REGION}:{_PG_HOST_INFO.url}:{expected_default_pg_port}:admin")
205204
assert _GENERATED_TOKEN == actual_token.token
206-
assert actual_token.is_expired() is False
205+
assert not actual_token.is_expired()
207206

208207
# Assert password has been updated to the value in token cache
209208
expected_props = {"user": "admin", "iam_default_port": "0"}
@@ -217,7 +216,7 @@ def test_pg_connect_with_invalid_port_and_no_host_port_fall_backs_to_host_port(
217216
@patch("aws_advanced_python_wrapper.iam_plugin.IamAuthPlugin._token_cache", _token_cache)
218217
def test_connect_expired_token_in_cache(user, mocker, mock_plugin_service, mock_session, mock_func, mock_client, mock_dialect):
219218
test_props: Properties = Properties({"user": user})
220-
cache_key, initial_token = set_token_cache(user, _PG_HOST_INFO.url, _DEFAULT_PG_PORT, _PG_REGION)
219+
cache_key, initial_token = set_token_cache(user, _PG_HOST_INFO.url, _DEFAULT_PG_PORT, _PG_REGION, True)
221220

222221
mock_func.side_effect = Exception("generic exception")
223222
target_plugin: IamAuthPlugin = IamAuthPlugin(mock_plugin_service,
@@ -234,7 +233,7 @@ def test_connect_expired_token_in_cache(user, mocker, mock_plugin_service, mock_
234233

235234
actual_token = _token_cache.get(cache_key)
236235
assert initial_token != actual_token
237-
assert actual_token.is_expired() is False
236+
assert not actual_token.is_expired()
238237

239238
if user == "admin":
240239
mock_client.generate_db_connect_admin_auth_token.assert_called_with(
@@ -280,7 +279,7 @@ def test_connect_empty_cache(user, mocker, mock_plugin_service, mock_connection,
280279
assert _GENERATED_TOKEN_NON_ADMIN == actual_token.token
281280

282281
assert mock_connection == actual_connection
283-
assert actual_token.is_expired() is False
282+
assert not actual_token.is_expired()
284283

285284

286285
@patch("aws_advanced_python_wrapper.iam_plugin.IamAuthPlugin._token_cache", _token_cache)
@@ -310,7 +309,7 @@ def test_connect_with_specified_port(mocker, mock_plugin_service, mock_session,
310309
assert _token_cache.get(_PG_CACHE_KEY) is None
311310
assert _GENERATED_TOKEN != actual_token.token
312311
assert f"{_TEST_TOKEN}:1234" == actual_token.token
313-
assert actual_token.is_expired() is False
312+
assert not actual_token.is_expired()
314313

315314
# Assert password has been updated to the value in token cache
316315
expected_props = {"user": "admin"}
@@ -346,7 +345,7 @@ def test_connect_with_specified_iam_default_port(mocker, mock_plugin_service, mo
346345
assert _token_cache.get(_PG_CACHE_KEY) is None
347346
assert _GENERATED_TOKEN != actual_token.token
348347
assert f"{_TEST_TOKEN}:{iam_default_port}" == actual_token.token
349-
assert actual_token.is_expired() is False
348+
assert not actual_token.is_expired()
350349

351350
# Assert password has been updated to the value in token cache
352351
expected_props = {"user": "admin", "iam_default_port": "9999"}
@@ -389,7 +388,7 @@ def test_connect_with_specified_region(user, mocker, mock_plugin_service, mock_s
389388

390389
expected_props = {"iam_region": "us-east-1", "user": user}
391390
actual_token = _token_cache.get(IamAuthUtils.get_cache_key(user, _PG_HOST_INFO.url, _DEFAULT_PG_PORT, iam_region))
392-
assert actual_token.is_expired() is False
391+
assert not actual_token.is_expired()
393392

394393
if user == "admin":
395394
mock_client.generate_db_connect_admin_auth_token.assert_called_with(
@@ -405,8 +404,8 @@ def test_connect_with_specified_region(user, mocker, mock_plugin_service, mock_s
405404

406405

407406
@pytest.mark.parametrize("iam_host", [
408-
pytest.param("dsqltestclusternamefoobar1.dsql-gamma.us-east-2.on.aws"),
409-
pytest.param("dsqltestclusternamefoobar2.dsql-gamma.us-east-2.on.aws"),
407+
pytest.param("dsqltestclusternamefoobar1.dsql.us-east-2.on.aws"),
408+
pytest.param("dsqltestclusternamefoobar2.dsql.us-east-2.on.aws"),
410409
])
411410
@patch("aws_advanced_python_wrapper.iam_plugin.IamAuthPlugin._token_cache", _token_cache)
412411
def test_connect_with_specified_host(iam_host: str, mocker, mock_plugin_service, mock_session, mock_func, mock_client, mock_dialect):
@@ -437,34 +436,9 @@ def test_connect_with_specified_host(iam_host: str, mocker, mock_plugin_service,
437436
assert actual_token is not None
438437
assert _GENERATED_TOKEN != actual_token.token
439438
assert f"{_TEST_TOKEN}:{iam_host}" == actual_token.token
440-
assert actual_token.is_expired() is False
439+
assert not actual_token.is_expired()
441440

442441

443442
def test_aws_supported_regions_url_exists():
444443
url = "https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/Concepts.RegionsAndAvailabilityZones.html"
445444
assert 200 == urllib.request.urlopen(url).getcode()
446-
447-
448-
@pytest.mark.parametrize("host", [
449-
pytest.param("<>"),
450-
pytest.param("#"),
451-
pytest.param("'"),
452-
pytest.param("\""),
453-
pytest.param("%"),
454-
pytest.param("^"),
455-
pytest.param("https://foo.com/abc.html"),
456-
pytest.param("foo.boo//"),
457-
pytest.param("8.8.8.8"),
458-
pytest.param("a.b"),
459-
])
460-
def test_invalid_iam_host(host, mocker, mock_plugin_service, mock_session, mock_func, mock_client, mock_dialect):
461-
test_props: Properties = Properties({"user": "admin"})
462-
with pytest.raises(AwsWrapperError):
463-
target_plugin: IamAuthPlugin = IamAuthPlugin(mock_plugin_service, DSQLTokenUtils(), mock_session)
464-
target_plugin.connect(
465-
target_driver_func=mocker.MagicMock(),
466-
driver_dialect=mock_dialect,
467-
host_info=HostInfo(host),
468-
props=test_props,
469-
is_initial_connection=False,
470-
connect_func=mock_func)

0 commit comments

Comments
 (0)