From 6258365df61aef6c4df498179b0eb39278c13f0d Mon Sep 17 00:00:00 2001
From: Alex Willmer <alex@moreati.org.uk>
Date: Wed, 9 Aug 2023 14:27:49 +0100
Subject: [PATCH] tests: Handle square bracket IPv6 in `docker port` output

Fixes
```
======================================================================
ERROR: setUpClass (ssh_test.BannerTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/vsts/work/1/s/tests/testlib.py", line 625, in setUpClass
    cls.dockerized_ssh = DockerizedSshDaemon(**daemon_args)
  File "/home/vsts/work/1/s/tests/testlib.py", line 553, in __init__
    self.start_container()
  File "/home/vsts/work/1/s/tests/testlib.py", line 533, in start_container
    self._get_container_port()
  File "/home/vsts/work/1/s/tests/testlib.py", line 510, in _get_container_port
    self.port = int(bport)
ValueError: invalid literal for int() with base 10: ':]:32770'
```
---
 tests/testlib.py | 29 ++++++++++++-----------------
 1 file changed, 12 insertions(+), 17 deletions(-)

diff --git a/tests/testlib.py b/tests/testlib.py
index f47d4cbc..242f211b 100644
--- a/tests/testlib.py
+++ b/tests/testlib.py
@@ -499,19 +499,18 @@ def get_docker_host():
 
 
 class DockerizedSshDaemon(object):
-    def _get_container_port(self):
-        s = subprocess.check_output(['docker', 'port', self.container_name])
-        for line in s.decode().splitlines():
-            m = self.PORT_RE.match(line)
-            if not m:
-                continue
-            dport, proto, _, bport = m.groups()
-            if dport == '22' and proto == 'tcp':
-                self.port = int(bport)
+    PORT_RE = re.compile(
+        # e.g. 0.0.0.0:32771, :::32771, [::]:32771'
+        r'(?P<addr>[0-9.]+|::|\[[a-f0-9:.]+\]):(?P<port>[0-9]+)',
+    )
 
-        self.host = self.get_host()
-        if self.port is None:
+    @classmethod
+    def get_port(cls, container):
+        s = subprocess.check_output(['docker', 'port', container, '22/tcp'])
+        m = cls.PORT_RE.search(s.decode())
+        if not m:
             raise ValueError('could not find SSH port in: %r' % (s,))
+        return int(m.group('port'))
 
     def start_container(self):
         try:
@@ -530,7 +529,6 @@ class DockerizedSshDaemon(object):
             self.image,
         ]
         subprocess.check_output(args)
-        self._get_container_port()
 
     def __init__(self, mitogen_test_distro=os.environ.get('MITOGEN_TEST_DISTRO', 'debian9')):
         if '-'  in mitogen_test_distro:
@@ -545,12 +543,9 @@ class DockerizedSshDaemon(object):
             self.python_path = '/usr/bin/python'
 
         self.image = 'public.ecr.aws/n5z0e8q9/%s-test' % (distro,)
-
-        # 22/tcp -> 0.0.0.0:32771
-        self.PORT_RE = re.compile(r'([^/]+)/([^ ]+) -> ([^:]+):(.*)')
-        self.port = None
-
         self.start_container()
+        self.host = self.get_host()
+        self.port = self.get_port(self.container_name)
 
     def get_host(self):
         return get_docker_host()