From 57839104045f839476536330a0cb53da8d586bc4 Mon Sep 17 00:00:00 2001
From: Kenneth Hoste <kenneth.hoste@ugent.be>
Date: Wed, 6 Mar 2024 12:03:03 +0100
Subject: [PATCH] add support to run_shell_cmd for providing list of answers in
 qa_patterns + fix output for hooks triggered for interactive shell commands
 run with run_shell_cmd

---
 easybuild/tools/run.py | 20 ++++++++++++++++++--
 test/framework/run.py  | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/easybuild/tools/run.py b/easybuild/tools/run.py
index c9a4b3d06f..ba40b704ac 100644
--- a/easybuild/tools/run.py
+++ b/easybuild/tools/run.py
@@ -319,7 +319,11 @@ def to_cmd_str(cmd):
 
     if with_hooks:
         hooks = load_hooks(build_option('hooks'))
-        hook_res = run_hook(RUN_SHELL_CMD, hooks, pre_step_hook=True, args=[cmd], kwargs={'work_dir': work_dir})
+        kwargs = {
+            'interactive': bool(qa_patterns),
+            'work_dir': work_dir,
+        }
+        hook_res = run_hook(RUN_SHELL_CMD, hooks, pre_step_hook=True, args=[cmd], kwargs=kwargs)
         if hook_res:
             cmd, old_cmd = hook_res, cmd
             cmd_str = to_cmd_str(cmd)
@@ -375,10 +379,21 @@ def to_cmd_str(cmd):
 
             # only consider answering questions if there's new output beyond additional whitespace
             if qa_patterns:
-                for question, answer in qa_patterns:
+                for question, answers in qa_patterns:
+
                     question += r'[\s\n]*$'
                     regex = re.compile(question.encode())
                     if regex.search(stdout):
+                        # if answer is specified as a list, we take the first item as current answer,
+                        # and add it to the back of the list (so we cycle through answers)
+                        if isinstance(answers, list):
+                            answer = answers.pop(0)
+                            answers.append(answer)
+                        elif isinstance(answers, str):
+                            answer = answers
+                        else:
+                            raise EasyBuildError(f"Unknown type of answers encountered: {answers}")
+
                         answer += '\n'
                         os.write(proc.stdin.fileno(), answer.encode())
                         time_no_match = 0
@@ -437,6 +452,7 @@ def to_cmd_str(cmd):
     if with_hooks:
         run_hook_kwargs = {
             'exit_code': res.exit_code,
+            'interactive': bool(qa_patterns),
             'output': res.output,
             'stderr': res.stderr,
             'work_dir': res.work_dir,
diff --git a/test/framework/run.py b/test/framework/run.py
index aa22ef88fd..6fb8d2b072 100644
--- a/test/framework/run.py
+++ b/test/framework/run.py
@@ -996,6 +996,27 @@ def test_run_cmd_qa_answers(self):
         self.assertEqual(out, "question\nanswer1\nquestion\nanswer2\n" * 2)
         self.assertEqual(ec, 0)
 
+    def test_run_shell_cmd_qa_answers(self):
+        """Test providing list of answers for a question in run_shell_cmd."""
+
+        cmd = "echo question; read x; echo $x; " * 2
+        qa = [("question", ["answer1", "answer2"])]
+
+        with self.mocked_stdout_stderr():
+            res = run_shell_cmd(cmd, qa_patterns=qa)
+        self.assertEqual(res.output, "question\nanswer1\nquestion\nanswer2\n")
+        self.assertEqual(res.exit_code, 0)
+
+        with self.mocked_stdout_stderr():
+            self.assertErrorRegex(EasyBuildError, "Unknown type of answers encountered", run_shell_cmd, cmd, qa_patterns=[('question', 1)])
+
+        # test cycling of answers
+        cmd = cmd * 2
+        with self.mocked_stdout_stderr():
+            res = run_shell_cmd(cmd, qa_patterns=qa)
+        self.assertEqual(res.output, "question\nanswer1\nquestion\nanswer2\n" * 2)
+        self.assertEqual(res.exit_code, 0)
+
     def test_run_cmd_simple(self):
         """Test return value for run_cmd in 'simple' mode."""
         with self.mocked_stdout_stderr():
@@ -1124,6 +1145,14 @@ def test_run_cmd_dry_run(self):
         self.assertEqual(read_file(outfile), "This is always echoed\n")
 
         # Q&A commands
+        self.mock_stdout(True)
+        run_shell_cmd("some_qa_cmd", qa_patterns=[('question1', 'answer1')])
+        stdout = self.get_stdout()
+        self.mock_stdout(False)
+
+        expected = """  running interactive shell command "some_qa_cmd"\n"""
+        self.assertIn(expected, stdout)
+
         self.mock_stdout(True)
         run_cmd_qa("some_qa_cmd", {'question1': 'answer1'})
         stdout = self.get_stdout()
@@ -1565,6 +1594,17 @@ def post_run_shell_cmd_hook(cmd, *args, **kwargs):
         ])
         self.assertEqual(stdout, expected_stdout)
 
+        with self.mocked_stdout_stderr():
+            run_shell_cmd("sleep 2; make", qa_patterns=[('q', 'a')])
+            stdout = self.get_stdout()
+
+        expected_stdout = '\n'.join([
+            "pre-run hook interactive 'sleep 2; make' in %s" % cwd,
+            "post-run hook interactive 'sleep 2; echo make' (exit code: 0, output: 'make\n')",
+            '',
+        ])
+        self.assertEqual(stdout, expected_stdout)
+
         with self.mocked_stdout_stderr():
             run_cmd_qa("sleep 2; make", qa={})
             stdout = self.get_stdout()