From 2e4f63e62f5c3ed54e2ea1bcadcdbb301a70bc5f Mon Sep 17 00:00:00 2001
From: Roland Dobai <roland@espressif.com>
Date: Mon, 6 Nov 2023 13:22:35 +0100
Subject: [PATCH] ci(fix): Improve the stability of the autocomplete tests

---
 tools/ci/test_autocomplete.py | 84 +++++++++++++++++------------------
 1 file changed, 40 insertions(+), 44 deletions(-)

diff --git a/tools/ci/test_autocomplete.py b/tools/ci/test_autocomplete.py
index 9ee4406dc9..36dfe00c22 100755
--- a/tools/ci/test_autocomplete.py
+++ b/tools/ci/test_autocomplete.py
@@ -1,56 +1,52 @@
 #!/usr/bin/env python
+# SPDX-FileCopyrightText: 2022-2023 Espressif Systems (Shanghai) CO LTD
+# SPDX-License-Identifier: Apache-2.0
+import collections
 import os
-import sys
 import unittest
 
 import pexpect
 
+Test = collections.namedtuple('Test', ['name', 'term', 'pattern', 'ext'])
 
-class Test(unittest.TestCase):
-    def test_fish(self):
-        os.environ['TERM'] = 'vt100'
-        child = pexpect.spawn('fish -i')
-        with open(os.environ['IDF_PATH'] + '/fish' + str(sys.version_info.major) + '.out', 'wb') as output:
-            child.logfile = output
-            child.sendline('. ./export.fish')
-            result = child.expect(
-                ['Go to the project directory and run.*idf\\.py build', pexpect.EOF,
-                 pexpect.TIMEOUT], timeout=40)
-            self.assertEqual(result, 0, 'Export was not successful!')
-            child.send('idf.py \t\t')
-            result = child.expect(['all.*app.*app-flash.*bootloader.*', pexpect.EOF, pexpect.TIMEOUT], timeout=40)
-            self.assertEqual(result, 0, 'Autocompletion for idf.py failed in fish!')
 
-    def test_bash(self):
-        os.environ['TERM'] = 'xterm-256color'
-        child = pexpect.spawn('bash -i')
-        with open(os.environ['IDF_PATH'] + '/bash' + str(sys.version_info.major) + '.out', 'wb') as output:
-            child.logfile = output
-            child.sendline('. ./export.sh')
-            child.send('idf.py \t\t')
-            result = child.expect(
-                ['Go to the project directory and run.*idf\\.py build', pexpect.EOF,
-                 pexpect.TIMEOUT], timeout=40)
-            self.assertEqual(result, 0, 'Export was not successful!')
-            result = child.expect(
-                ['all.*app.*app-flash.*bootloader.*bootloader-flash.*build-system-targets.*clean.*', pexpect.EOF,
-                 pexpect.TIMEOUT], timeout=40)
-            self.assertEqual(result, 0, 'Autocompletion for idf.py failed in bash!')
+TESTS = (Test('fish', 'vt100', 'all.*app.*app-flash.*bootloader.*', 'fish'),
+         Test('bash', 'xterm-256color', 'all.*app.*app-flash.*bootloader.*bootloader-flash.*build-system-targets.*clean.*', 'sh'),
+         Test('zsh', '', 'all.*app.*app-flash.*bootloader.*bootloader-flash.*build-system-targets.*clean.*', 'sh'))
 
-    def test_zsh(self):
-        child = pexpect.spawn('zsh -i')
-        with open(os.environ['IDF_PATH'] + '/zsh' + str(sys.version_info.major) + '.out', 'wb') as output:
-            child.logfile = output
-            child.sendline('. ./export.sh')
-            result = child.expect(
-                ['Go to the project directory and run.*idf\\.py build', pexpect.EOF,
-                 pexpect.TIMEOUT], timeout=40)
-            self.assertEqual(result, 0, 'Export was not successful!')
-            child.send('idf.py \t\t')
-            result = child.expect(
-                ['all.*app.*app-flash.*bootloader.*bootloader-flash.*build-system-targets.*clean.*', pexpect.EOF,
-                 pexpect.TIMEOUT], timeout=40)
-            self.assertEqual(result, 0, 'Autocompletion for idf.py failed in zsh!')
+
+# Additional positional arguments for all child.expect() calls are constant so we can rely on the order and print message
+# about which pattern was matched
+pargs = (pexpect.EOF, pexpect.TIMEOUT)
+
+
+def get_fail_msg(pproc, msg, index):
+    try:
+        buf = pproc._buffer.getvalue()
+    except AttributeError:
+        # _buffer is an internal of pexpect.spawn and is not part of the API.
+        # Either there is no _buffer or it is not io.BytesIO() anymore.
+        buf = '<UNKNOWN>'
+
+    return '{} ({}) buffer: "{}"'.format(msg, 'EOF - child has exited' if index == 1 else 'TIMEOUT', buf)
+
+
+class UTTest(unittest.TestCase):
+
+    def test_shell(self):
+        idf_path = os.environ['IDF_PATH']
+        env = os.environ.copy()
+        for test in TESTS:
+            with self.subTest():
+                with open(os.path.join(idf_path, f'{test.name}.out'), 'wb') as o:
+                    env['TERM'] = test.term
+                    with pexpect.spawn(f'{test.name} -i', env=env, logfile=o, timeout=200) as pproc:
+                        pproc.sendline(f'. {idf_path}/export.{test.ext}')
+                        i = pproc.expect(['Go to the project directory and run.*idf\\.py build', *pargs])
+                        self.assertEqual(i, 0, get_fail_msg(pproc, 'Export was not successful!', i))
+                        pproc.send('idf.py \t\t')
+                        i = pproc.expect([test.pattern, *pargs], timeout=100)
+                        self.assertEqual(i, 0, get_fail_msg(pproc, f'Autocompletion for idf.py failed in {test.name}!', i))
 
 
 if __name__ == '__main__':