diff --git a/patroni/ctl.py b/patroni/ctl.py index 83197bfa..8d9b21f4 100644 --- a/patroni/ctl.py +++ b/patroni/ctl.py @@ -60,6 +60,18 @@ class PatronictlPrettyTable(PrettyTable): self.__hline_num = 0 self.__hline = None + def __build_header(self, line): + header = self.__table_header[:len(line) - 2] + return "".join([line[0], header, line[1 + len(header):]]) + + def _stringify_hrule(self, *args, **kwargs): + ret = super(PatronictlPrettyTable, self)._stringify_hrule(*args, **kwargs) + where = args[1] if len(args) > 1 else kwargs.get('where') + if where == 'top_' and self.__table_header: + ret = self.__build_header(ret) + self.__hline_num += 1 + return ret + def _is_first_hline(self): return self.__hline_num == 0 @@ -71,8 +83,7 @@ class PatronictlPrettyTable(PrettyTable): # Inject nice table header if self._is_first_hline() and self.__table_header: - header = self.__table_header[:len(ret) - 2] - ret = "".join([ret[0], header, ret[1 + len(header):]]) + ret = self.__build_header(ret) self.__hline_num += 1 return ret diff --git a/tests/test_ctl.py b/tests/test_ctl.py index ed6e68a1..cfd75b39 100644 --- a/tests/test_ctl.py +++ b/tests/test_ctl.py @@ -7,10 +7,11 @@ from datetime import datetime, timedelta from mock import patch, Mock from patroni.ctl import ctl, store_config, load_config, output_members, get_dcs, parse_dcs, \ get_all_members, get_any_member, get_cursor, query_member, configure, PatroniCtlException, apply_config_changes, \ - format_config_for_editing, show_diff, invoke_editor, format_pg_version, CONFIG_FILE_PATH + format_config_for_editing, show_diff, invoke_editor, format_pg_version, CONFIG_FILE_PATH, PatronictlPrettyTable from patroni.dcs.etcd import AbstractEtcdClientWithFailover, Failover from patroni.psycopg import OperationalError from patroni.utils import tzutc +from prettytable import PrettyTable, ALL from urllib3 import PoolManager from . import MockConnect, MockCursor, MockResponse, psycopg_connect @@ -648,3 +649,24 @@ class TestCtl(unittest.TestCase): result = self.runner.invoke(ctl, ['reinit', 'alpha', 'other', '--wait'], input='y\ny') self.assertIn("Waiting for reinitialize to complete on: other", result.output) self.assertIn("Reinitialize is completed on: other", result.output) + + +class TestPatronictlPrettyTable(unittest.TestCase): + + def setUp(self): + self.pt = PatronictlPrettyTable(' header', ['foo', 'bar'], hrules=ALL) + + def test__get_hline(self): + expected = '+-----+-----+' + self.pt._hrule = expected + self.assertEqual(self.pt._hrule, '+ header----+') + self.assertFalse(self.pt._is_first_hline()) + self.assertEqual(self.pt._hrule, expected) + + @patch.object(PrettyTable, '_stringify_hrule', Mock(return_value='+-----+-----+')) + def test__stringify_hrule(self): + self.assertEqual(self.pt._stringify_hrule((), 'top_'), '+ header----+') + self.assertFalse(self.pt._is_first_hline()) + + def test_output(self): + self.assertEqual(str(self.pt), '+ header----+\n| foo | bar |\n+-----+-----+')