From d17b16c9b1bf3e325d415b3777c2a7bd24c764ce Mon Sep 17 00:00:00 2001
From: Aleksei Zotov <azotcsit@gmail.com>
Date: Fri, 18 Feb 2022 20:37:18 +0400
Subject: [PATCH] Fix flaky test - test_cqlsh_completion.TestCqlshCompletion

Patch by Aleksei Zotov; reviewed by Brandon Williams for CASSANDRA-17338
---
 CHANGES.txt                                  |  1 +
 pylib/cqlshlib/test/run_cqlsh.py             |  4 ++--
 pylib/cqlshlib/test/test_cqlsh_completion.py | 19 ++++++++++++++++---
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/CHANGES.txt b/CHANGES.txt
index 6705934aa0..40e30f20ab 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -1,4 +1,5 @@
 3.0.27
+ * Fix flaky test - test_cqlsh_completion.TestCqlshCompletion (CASSANDRA-17338)
  * Fixed TestCqlshOutput failing tests (CASSANDRA-17386)
  * LeveledCompactionStrategy disk space check improvements (CASSANDRA-17272)
  * Lazy transaction log replica creation allows incorrect replica content divergence during anticompaction (CASSANDRA-17273)
diff --git a/pylib/cqlshlib/test/run_cqlsh.py b/pylib/cqlshlib/test/run_cqlsh.py
index b011df4acd..3db673bf16 100644
--- a/pylib/cqlshlib/test/run_cqlsh.py
+++ b/pylib/cqlshlib/test/run_cqlsh.py
@@ -268,8 +268,8 @@ class CqlshRunner(ProcRunner):
         else:
             self.output_header = self.read_to_next_prompt()
 
-    def read_to_next_prompt(self):
-        return self.read_until(self.prompt, timeout=10.0, ptty_timeout=3)
+    def read_to_next_prompt(self, timeout=10.0):
+        return self.read_until(self.prompt, timeout=timeout, ptty_timeout=3)
 
     def read_up_to_timeout(self, timeout, blksize=4096):
         output = ProcRunner.read_up_to_timeout(self, timeout, blksize=blksize)
diff --git a/pylib/cqlshlib/test/test_cqlsh_completion.py b/pylib/cqlshlib/test/test_cqlsh_completion.py
index 75198b6867..87315bcd62 100644
--- a/pylib/cqlshlib/test/test_cqlsh_completion.py
+++ b/pylib/cqlshlib/test/test_cqlsh_completion.py
@@ -19,9 +19,12 @@
 
 from __future__ import with_statement
 
+import locale
+import os
 import re
 from .basecase import BaseTestCase, cqlsh
 from .cassconnect import testrun_cqlsh
+from .run_cqlsh import TimeoutError
 import unittest
 import sys
 
@@ -42,7 +45,11 @@ completion_separation_re = re.compile(r'\s+')
 class CqlshCompletionCase(BaseTestCase):
 
     def setUp(self):
-        self.cqlsh_runner = testrun_cqlsh(cqlver=None, env={'COLUMNS': '100000'})
+        env = os.environ.copy()
+        env['COLUMNS'] = '100000'
+        if (locale.getpreferredencoding() != 'UTF-8'):
+             env['LC_CTYPE'] = 'en_US.utf8'
+        self.cqlsh_runner = testrun_cqlsh(cqlver=None, env=env)
         self.cqlsh = self.cqlsh_runner.__enter__()
 
     def tearDown(self):
@@ -132,8 +139,14 @@ class CqlshCompletionCase(BaseTestCase):
                                        other_choices_ok=other_choices_ok,
                                        split_completed_lines=split_completed_lines)
         finally:
-            self.cqlsh.send(CTRL_C)  # cancel any current line
-            self.cqlsh.read_to_next_prompt()
+            try:
+                self.cqlsh.send(CTRL_C)  # cancel any current line
+                self.cqlsh.read_to_next_prompt(timeout=1.0)
+            except TimeoutError:
+                # retry once
+                self.cqlsh.send(CTRL_C)
+                self.cqlsh.read_to_next_prompt(timeout=10.0)
+
 
     def strategies(self):
         return self.module.CqlRuleSet.replication_strategies