diff --git a/aws_gate/decorators.py b/aws_gate/decorators.py index 1f702f1f..984b7e27 100644 --- a/aws_gate/decorators.py +++ b/aws_gate/decorators.py @@ -2,6 +2,7 @@ import os from subprocess import PIPE +from packaging.version import parse as parse_version from wrapt import decorator from aws_gate.constants import PLUGIN_INSTALL_PATH, PLUGIN_NAME @@ -43,9 +44,7 @@ def wrapper( required_version, ) - if version and int(version.replace(".", "")) < int( - required_version.replace(".", "") - ): + if version and parse_version(version) < parse_version(required_version): raise ValueError("Invalid plugin version: {}".format(version)) return wrapped_function(*args, **kwargs) diff --git a/requirements/requirements.txt b/requirements/requirements.txt index aba8563a..1dd6bbf0 100644 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -1,7 +1,8 @@ boto3~=1.15.15 +cryptography==3.1.1 marshmallow==3.8.0 +packaging==20.4 PyYAML>=5.1,<5.4 requests==2.24.0 -cryptography==3.1.1 -wrapt==1.12.1 unix-ar==0.2.1 +wrapt==1.12.1 diff --git a/tests/unit/test_decorators.py b/tests/unit/test_decorators.py index 60acf075..7437fbc6 100644 --- a/tests/unit/test_decorators.py +++ b/tests/unit/test_decorators.py @@ -41,8 +41,9 @@ def test_function(): test_function() -def test_plugin_version(mocker): - m = mocker.patch("aws_gate.decorators.execute_plugin", return_value="1.1.23.0") +@pytest.mark.parametrize("version", ["1.1.23.0", "1.2.7.0"]) +def test_plugin_version(mocker, version): + m = mocker.patch("aws_gate.decorators.execute_plugin", return_value=version) @plugin_version("1.1.23.0") def test_function(): @@ -52,10 +53,11 @@ def test_function(): assert m.call_args == mocker.call(["--version"], stdout=PIPE, stderr=PIPE) -def test_plugin_version_bad_version(mocker): +@pytest.mark.parametrize("version", ["1.1.25.0", "1.2.7.0"]) +def test_plugin_version_bad_version(mocker, version): mocker.patch("aws_gate.decorators.execute_plugin", return_value="1.1.23.0") - @plugin_version("1.1.25.0") + @plugin_version(version) def test_function(): return "executed"