fix: use regex for go versions

Author: sriram-mv

aws_lambda_builders/workflows/go_modules/validator.py

@@ -3,6 +3,7 @@
 """
 
 import logging
+import re
 import os
 import subprocess
 
@@ -12,9 +13,9 @@
 
 
 class GoRuntimeValidator(object):
-
     LANGUAGE = "go"
     SUPPORTED_RUNTIMES = {"go1.x"}
+    GO_VERSION_REGEX = re.compile("go(\\d)\\.(x|\\d+)")
 
     def __init__(self, runtime):
         self.runtime = runtime
@@ -29,13 +30,13 @@ def has_runtime(self):
         return self.runtime in self.SUPPORTED_RUNTIMES
 
     @staticmethod
-    def sanitize_go_version_part(version_part):
-        sanitized_version_part = ""
-        for char in version_part:
-            if char.isalpha():
-                break
-            sanitized_version_part += char
-        return sanitized_version_part
+    def get_go_versions(version_string):
+        parts = GoRuntimeValidator.GO_VERSION_REGEX.findall(version_string)
+        try:
+            # NOTE(sriram-mv): The version parts need to be a list with a major and minor version.
+            return int(parts[0][0]), int(parts[0][1])
+        except IndexError:
+            return 0, 0
 
     def validate(self, runtime_path):
         """
@@ -51,19 +52,13 @@ def validate(self, runtime_path):
         min_expected_minor_version = 11 if expected_major_version == 1 else 0
 
         p = subprocess.Popen([runtime_path, "version"], cwd=os.getcwd(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-        out, _ = p.communicate()
+        version_string, _ = p.communicate()
 
         if p.returncode == 0:
-            out_parts = out.decode().split()
-            if len(out_parts) >= 3:
-                version_parts = [
-                    int(GoRuntimeValidator.sanitize_go_version_part(version_part))
-                    for version_part in out_parts[2].replace(self.LANGUAGE, "").split(".")
-                ]
-                if len(version_parts) >= 2:
-                    if version_parts[0] == expected_major_version and version_parts[1] >= min_expected_minor_version:
-                        self._valid_runtime_path = runtime_path
-                        return self._valid_runtime_path
+            major_version, minor_version = GoRuntimeValidator.get_go_versions(version_string.decode())
+            if major_version == expected_major_version and minor_version >= min_expected_minor_version:
+                self._valid_runtime_path = runtime_path
+                return self._valid_runtime_path
 
         # otherwise, raise mismatch exception
         raise MisMatchRuntimeError(language=self.LANGUAGE, required_runtime=self.runtime, runtime_path=runtime_path)

tests/unit/workflows/go_modules/test_validator.py

@@ -14,7 +14,7 @@ def __init__(self, returncode, out=b"", err=b""):
         self.err = err
 
     def communicate(self):
-        return (self.out, self.err)
+        return self.out, self.err
 
 
 class TestGoRuntimeValidator(TestCase):
@@ -30,38 +30,50 @@ def test_runtime_validate_unsupported_language_fail_open(self):
         validator = GoRuntimeValidator(runtime="go2.x")
         validator.validate(runtime_path="/usr/bin/go2")
 
+    @parameterized.expand(
+        [
+            ("go1.11.2", (1, 11)),
+            ("go1.11rc.2", (1, 11)),
+            ("go1.16beta1", (1, 16)),
+            ("go%$", (0, 0)),
+            ("unknown", (0, 0)),
+        ]
+    )
+    def test_get_go_versions(self, version_string, version_parts):
+        self.assertEqual(self.validator.get_go_versions(version_string), version_parts)
+
     @parameterized.expand(
         [(b"go version go1.11.2 test",), (b"go version go1.11rc.2 test",), (b"go version go1.16beta1 test",)]
     )
     def test_runtime_validate_supported_version_runtime(self, go_version_output):
         with mock.patch("subprocess.Popen") as mock_subprocess:
             mock_subprocess.return_value = MockSubProcess(0, out=go_version_output)
             self.validator.validate(runtime_path="/usr/bin/go")
-            self.assertTrue(mock_subprocess.call_count, 1)
+            self.assertEqual(mock_subprocess.call_count, 1)
 
     def test_runtime_validate_supported_higher_than_min_version_runtime(self):
         with mock.patch("subprocess.Popen") as mock_subprocess:
             mock_subprocess.return_value = MockSubProcess(0, out=b"go version go1.12 test")
             self.validator.validate(runtime_path="/usr/bin/go")
-            self.assertTrue(mock_subprocess.call_count, 1)
+            self.assertEqual(mock_subprocess.call_count, 1)
 
     def test_runtime_validate_mismatch_nonzero_exit(self):
         with mock.patch("subprocess.Popen") as mock_subprocess:
             mock_subprocess.return_value = MockSubProcess(1)
             with self.assertRaises(MisMatchRuntimeError):
                 self.validator.validate(runtime_path="/usr/bin/go")
-                self.assertTrue(mock_subprocess.call_count, 1)
+                self.assertEqual(mock_subprocess.call_count, 1)
 
     def test_runtime_validate_mismatch_invalid_version(self):
         with mock.patch("subprocess.Popen") as mock_subprocess:
             mock_subprocess.return_value = MockSubProcess(0, out=b"go version")
             with self.assertRaises(MisMatchRuntimeError):
                 self.validator.validate(runtime_path="/usr/bin/go")
-                self.assertTrue(mock_subprocess.call_count, 1)
+                self.assertEqual(mock_subprocess.call_count, 1)
 
     def test_runtime_validate_mismatch_minor_version(self):
         with mock.patch("subprocess.Popen") as mock_subprocess:
             mock_subprocess.return_value = MockSubProcess(0, out=b"go version go1.10.2 test")
             with self.assertRaises(MisMatchRuntimeError):
                 self.validator.validate(runtime_path="/usr/bin/go")
-                self.assertTrue(mock_subprocess.call_count, 1)
+                self.assertEqual(mock_subprocess.call_count, 1)