tests/termex.py

   1 #!/usr/bin/env python3
   2
   3 # Copyright (C) 2016-2017 Petr Pavlu <setup@dagobah.cz>
   4 #
   5 # This file is part of CenterIM.
   6 #
   7 # CenterIM is free software; you can redistribute it and/or modify
   8 # it under the terms of the GNU General Public License as published by
   9 # the Free Software Foundation; either version 2 of the License, or
  10 # (at your option) any later version.
  11 #
  12 # CenterIM is distributed in the hope that it will be useful,
  13 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 # GNU General Public License for more details.
  16 #
  17 # You should have received a copy of the GNU General Public License
  18 # along with CenterIM.  If not, see <http://www.gnu.org/licenses/>.
  19
  20 """Termex terminal emulator and test framework."""
  21
  22 import argparse
  23 import difflib
  24 import errno
  25 import os
  26 import pty
  27 import re
  28 import selectors
  29 import sys
  30 import time
  31 import tkinter
  32 import xml.etree.ElementTree as ElementTree
  33
  34 # The module relies on selectors.select() to automatically retry the operation
  35 # with a recomputed timeout when it gets interrupted by a signal. This
  36 # behaviour was introduced in Python 3.5.
  37 if sys.hexversion < 0x03050000:
  38     print("This program requires at least Python 3.5.", file=sys.stderr)
  39     sys.exit(1)
  40
  41 ROWS = 24
  42 COLUMNS = 80
  43 CHILD_TIMEOUT = 5
  44
  45 ATTR_NORMAL = 0
  46 ATTR_REVERSE = 1
  47
  48 COLOR_BLACK = 0
  49 COLOR_RED = 1
  50 COLOR_GREEN = 2
  51 COLOR_YELLOW = 3
  52 COLOR_BLUE = 4
  53 COLOR_MAGENTA = 5
  54 COLOR_CYAN = 6
  55 COLOR_WHITE = 7
  56 COLOR_DEFAULT = 9
  57 COLOR_REGISTER = ((COLOR_BLACK, 'black'), (COLOR_RED, 'red'),
  58                   (COLOR_GREEN, 'green'), (COLOR_YELLOW, 'yellow'),
  59                   (COLOR_BLUE, 'blue'), (COLOR_MAGENTA, 'magenta'),
  60                   (COLOR_CYAN, 'cyan'), (COLOR_WHITE, 'white'),
  61                   (COLOR_DEFAULT, 'default'))
  62 COLOR_TO_STRING_MAP = {id_: name for id_, name in COLOR_REGISTER}
  63 STRING_TO_COLOR_MAP = {name: id_ for id_, name in COLOR_REGISTER}
  64 COLORS = {id_ for id_, _ in COLOR_REGISTER}
  65 REAL_COLOR_NAMES = tuple(
  66     name for id_, name in COLOR_REGISTER if id_ != COLOR_DEFAULT)
  67
  68 COLOR_DEFAULT_FOREGROUND = COLOR_BLACK
  69 COLOR_DEFAULT_BACKGROUND = COLOR_WHITE
  70
  71 CODE_ENTER = '\x0d'
  72 CODE_FN = ('\x1bOP', '\x1bOQ', '\x1bOR', '\x1bOS', '\x1b[15~', '\x1b[17~',
  73            '\x1b[18~', '\x1b[19~', '\x1b[20~', '\x1b[21~', '\x1b[23~')
  74 CODE_PAGE_UP = '\x1b[5~'
  75 CODE_PAGE_DOWN = '\x1b[6~'
  76
  77
  78 def attr_to_string(attr):
  79     """Get string representation of given attributes."""
  80
  81     res = []
  82     if (attr & ATTR_REVERSE) != 0:
  83         res.append('reverse')
  84     return '|'.join(res)
  85
  86
  87 def string_to_attr(string):
  88     """
  89     Convert a string to attributes. Exception ValueError is raised if some
  90     attribute is invalid.
  91     """
  92
  93     res = ATTR_NORMAL
  94     for attr in string.split('|'):
  95         if attr == 'normal':
  96             pass
  97         elif attr == 'reverse':
  98             res |= ATTR_REVERSE
  99         else:
 100             raise ValueError("Unrecognized attribute '{}'".format(attr))
 101     return res
 102
 103
 104 def color_to_string(color):
 105     """Get string representation of a given color."""
 106
 107     return COLOR_TO_STRING_MAP[color]
 108
 109
 110 def string_to_color(string):
 111     """
 112     Convert a string to a color. Exception ValueError is raised if the color
 113     name is not recognized.
 114     """
 115
 116     try:
 117         return STRING_TO_COLOR_MAP[string]
 118     except KeyError:
 119         raise ValueError("Unrecognized color '{}'".format(string))
 120
 121
 122 class TermChar:
 123     """On-screen character."""
 124
 125     def __init__(self, char=" ", attr=ATTR_NORMAL, fgcolor=COLOR_DEFAULT,
 126                  bgcolor=COLOR_DEFAULT):
 127         self.char = char
 128         self.attr = attr
 129         self.fgcolor = fgcolor
 130         self.bgcolor = bgcolor
 131
 132     def __eq__(self, other):
 133         return self.__dict__ == other.__dict__
 134
 135     def _get_translated_fgcolor(self):
 136         """
 137         Return the foreground color. If the current color is COLOR_DEFAULT then
 138         COLOR_DEFAULT_FOREGROUND is returned.
 139         """
 140
 141         if self.fgcolor == COLOR_DEFAULT:
 142             return COLOR_DEFAULT_FOREGROUND
 143         return self.fgcolor
 144
 145     def _get_translated_bgcolor(self):
 146         """
 147         Return the background color. If the current color is COLOR_DEFAULT then
 148         COLOR_DEFAULT_BACKGROUND is returned.
 149         """
 150
 151         if self.bgcolor == COLOR_DEFAULT:
 152             return COLOR_DEFAULT_BACKGROUND
 153         return self.bgcolor
 154
 155     def get_tag_foreground(self):
 156         """
 157         Return a name of the final foreground color that should be used to
 158         display the character on the screen.
 159         """
 160
 161         if self.attr & ATTR_REVERSE:
 162             color = self._get_translated_bgcolor()
 163         else:
 164             color = self._get_translated_fgcolor()
 165         return color_to_string(color)
 166
 167     def get_tag_background(self):
 168         """
 169         Return a name of the final background color that should be used to
 170         display the character on the screen.
 171         """
 172
 173         if self.attr & ATTR_REVERSE:
 174             color = self._get_translated_fgcolor()
 175         else:
 176             color = self._get_translated_bgcolor()
 177         return color_to_string(color)
 178
 179
 180 class Term:
 181     """Termex terminal emulator."""
 182
 183     MODE_RUN = 0
 184     MODE_RECORD = 1
 185     MODE_TEST = 2
 186
 187     class _TerminalConnectionException(Exception):
 188         """
 189         Exception reported when communication with the pseudo-terminal fails.
 190         """
 191         pass
 192
 193     def __init__(self, root, program, mode, terminfo=None):
 194         self._root = root
 195         self._program = program
 196         self._mode = mode
 197         self._terminfo = terminfo
 198
 199         # Test mode obtains the program name from the playbook.
 200         if self._mode == self.MODE_TEST:
 201             assert self._program is None
 202
 203         self._child_pid = None
 204         self._fd = None
 205
 206         self._screen = None
 207         self._cur_y = 0
 208         self._cur_x = 0
 209         self._attr = ATTR_NORMAL
 210         self._fgcolor = COLOR_DEFAULT
 211         self._bgcolor = COLOR_DEFAULT
 212         self._charbuf = b''
 213
 214         # Initialize the GUI if requested.
 215         if self._root:
 216             self._root.title("Termex")
 217             self._frame = tkinter.Frame(self._root)
 218
 219             self._text = tkinter.Text(self._root, height=ROWS, width=COLUMNS)
 220             self._text.config(
 221                 foreground=color_to_string(COLOR_DEFAULT_FOREGROUND),
 222                 background=color_to_string(COLOR_DEFAULT_BACKGROUND))
 223             self._text.pack()
 224
 225             # Configure tag values.
 226             for fgcolor_str in REAL_COLOR_NAMES:
 227                 for bgcolor_str in REAL_COLOR_NAMES:
 228                     tag = 'tag_{}-{}'.format(fgcolor_str, bgcolor_str)
 229                     self._text.tag_config(tag, foreground=fgcolor_str,
 230                                           background=bgcolor_str)
 231
 232         self._erase_all()
 233
 234         if self._mode == self.MODE_RECORD:
 235             self._test_e = ElementTree.Element('test')
 236             self._test_e.set('program', self._program)
 237
 238     def _start_program(self):
 239         """
 240         Fork, connect the child's controlling terminal to a pseudo-terminal and
 241         start the selected child program.
 242
 243         Parent behaviour: Returns True when the fork was successful, False
 244         otherwise. Note that the returned value does not provide information
 245         whether the exec call in the child process was successful or not. That
 246         must be determined by attempting communication with the child.
 247
 248         Child behaviour: Execs the selected program and does not return if the
 249         call was successful, returns False otherwise.
 250         """
 251
 252         # Fork and connect the child's controlling terminal to a
 253         # pseudo-terminal.
 254         try:
 255             self._child_pid, self._fd = pty.fork()
 256         except OSError as e:
 257             print("Fork to run '{}' failed: {}".format(self._program, e),
 258                   file=sys.stderr)
 259             return False
 260         if self._child_pid == 0:
 261             try:
 262                 env = {'PATH': '/bin:/usr/bin', 'TERM': 'termex',
 263                        'LC_ALL': 'en_US.UTF-8'}
 264                 if self._terminfo:
 265                     env['TERMINFO'] = self._terminfo
 266                 elif 'TERMINFO' in os.environ:
 267                     env['TERMINFO'] = os.environ['TERMINFO']
 268                 os.execle(self._program, self._program, env)
 269             except OSError as e:
 270                 print("Failed to execute '{}': {}".format(self._program, e),
 271                       file=sys.stderr)
 272                 return False
 273
 274         return True
 275
 276     def _finalize_program(self):
 277         """
 278         Close the connection to the pseudo-terminal and wait for the child
 279         program to complete. Returns True when the connection was successfully
 280         closed and the child completed in the timeout limit, False otherwise.
 281         """
 282
 283         res = True
 284
 285         # Close the file descriptor that is connected to the child's
 286         # controlling terminal.
 287         try:
 288             os.close(self._fd)
 289         except OSError as e:
 290             print("Failed to close file descriptor '{}' that is connected to "
 291                   "the child's controlling terminal: {}.".format(self._fd, e),
 292                   file=sys.stderr)
 293             res = False
 294
 295         # Wait for the child to finish. It should terminate now that its input
 296         # was closed.
 297         for _ in range(CHILD_TIMEOUT):
 298             try:
 299                 pid, _status = os.waitpid(self._child_pid, os.WNOHANG)
 300             except OSError as e:
 301                 print("Failed to wait on child '{}' to complete: "
 302                       "{}.".format(pid, e), file=sys.stderr)
 303                 res = False
 304                 break
 305             if pid != 0:
 306                 break
 307             time.sleep(1)
 308         else:
 309             print("Child '{}' has not completed.".format(self._child_pid),
 310                   file=sys.stderr)
 311             res = False
 312
 313         return res
 314
 315     def run_gui_mainloop(self):
 316         """Start the selected child program and run the tkinter's main loop."""
 317
 318         assert self._mode == self.MODE_RUN or self._mode == self.MODE_RECORD
 319
 320         # Start the specified program.
 321         if not self._start_program():
 322             return
 323
 324         try:
 325             # Prepare for running the main loop.
 326             self._root.createfilehandler(
 327                 self._fd, tkinter.READABLE,
 328                 lambda fd, mask: self._pty_callback())
 329             self._root.bind('<Key>', self._tk_key)
 330             self._root.bind('<<Quit>>', lambda e: self._quit_gui_mainloop())
 331             self._root.protocol('WM_DELETE_WINDOW', self._quit_gui_mainloop)
 332
 333             # Run the main loop.
 334             try:
 335                 self._root.mainloop()
 336             except self._TerminalConnectionException as e:
 337                 print("{}.".format(e), file=sys.stderr)
 338
 339             self._root.deletefilehandler(self._fd)
 340         finally:
 341             # Finalize the run of the child program.
 342             self._finalize_program()
 343
 344     def _quit_gui_mainloop(self):
 345         """Exit the tkinter's main loop."""
 346
 347         assert self._mode == self.MODE_RUN or self._mode == self.MODE_RECORD
 348         self._root.quit()
 349
 350     def _pty_callback(self):
 351         """
 352         Process a data event from the pseudo-terminal. Returns True when the
 353         connection to the pseudo-terminal was closed, False otherwise.
 354         Exception _TerminalConnectionException is raised if the read of the new
 355         data from the pseudo-terminal fails.
 356         """
 357
 358         closed = False
 359         try:
 360             char = os.read(self._fd, 1)
 361         except OSError as e:
 362             if e.errno == errno.EIO:
 363                 closed = True
 364             else:
 365                 raise self._TerminalConnectionException(
 366                     "Error reading from file descriptor '{}' that is "
 367                     "connected to the child's controlling terminal: "
 368                     "{}".format(self._fd, e))
 369
 370         # Check whether the descriptor referring to the pseudo-terminal slave
 371         # has been closed or end of file was reached.
 372         if closed or len(char) == 0:
 373             if self._root:
 374                 self._root.quit()
 375             return True
 376
 377         self._charbuf += char
 378
 379         if self._handle_sequence(self._charbuf):
 380             self._charbuf = b''
 381         else:
 382             # print("Unmatched {}.".format(self._charbuf), file=sys.stderr)
 383             pass
 384         return False
 385
 386     def _send_key(self, chars, name):
 387         """
 388         Write the specified characters that represent one key to the
 389         pseudo-terminal. If the recording mode is enabled then the specified
 390         key name is recorded in the test playbook. Exception
 391         _TerminalConnectionException is raised if the write to the
 392         pseudo-terminal fails.
 393         """
 394
 395         if self._mode == self.MODE_RECORD:
 396             # Record the key.
 397             action_e = ElementTree.SubElement(self._test_e, 'action')
 398             action_e.set('key', name)
 399             print("Recorded key '{}'.".format(name))
 400
 401         # Send the key to the terminal.
 402         try:
 403             os.write(self._fd, str.encode(chars))
 404         except OSError as e:
 405             raise self._TerminalConnectionException(
 406                 "Error writing characters '{}' to file descriptor '{}' that "
 407                 "is connected to the child's controlling terminal: "
 408                 "{}".format(chars, self._fd, e))
 409
 410     def _handle_sequence(self, seq):
 411         """
 412         Process a byte sequence received from the pseudo-terminal. Returns True
 413         when the sequence was recognized and successfully handled, False
 414         otherwise.
 415         """
 416
 417         if re.fullmatch(b'[^\x01-\x1f]+', seq):
 418             try:
 419                 uchar = seq.decode('utf-8')
 420                 self._print_char(
 421                     TermChar(uchar, self._attr, self._fgcolor, self._bgcolor))
 422                 return True
 423             except UnicodeError:
 424                 # Continue on the assumption that it is not yet a complete
 425                 # character. This assumption is wrong if the received text is
 426                 # actually malformed.
 427                 return False
 428
 429         if seq == b'\x07':
 430             # Bell.
 431             if self._root:
 432                 self._root.bell()
 433             return True
 434         if seq == b'\x08':
 435             # Backspace non-destructively.
 436             self._cur_x -= 1
 437             return True
 438         if seq == b'\x0d':
 439             # Go to beginning of line.
 440             self._cur_x = 0
 441             return True
 442         if seq == b'\x0a':
 443             # Move cursor down one line.
 444             self._cursor_down()
 445             return True
 446
 447         # Controls beginning with ESC.
 448
 449         # Control sequences.
 450         match = re.fullmatch(b'\x1b\\[([0-9]+)@', seq)
 451         if match:
 452             # Insert blank characters.
 453             self._insert_blanks(int(match.group(1)))
 454             return True
 455         if seq == b'\x1b[H':
 456             # Set cursor position to the default (top left).
 457             self._cur_y = 0
 458             self._cur_x = 0
 459             return True
 460         match = re.fullmatch(b'\x1b\\[([0-9]+);([0-9]+)H', seq)
 461         if match:
 462             # Set cursor position to (y,x).
 463             self._cur_y = int(match.group(1))
 464             self._cur_x = int(match.group(2))
 465             return True
 466         if self._charbuf == b'\x1b[K':
 467             # Erase in line to right.
 468             for x in range(self._cur_x, COLUMNS):
 469                 self._print_char_at(self._cur_y, x, TermChar())
 470             return True
 471         if seq == b'\x1b[2J':
 472             # Erase display completely.
 473             self._erase_all()
 474             return True
 475         if seq == b'\x1b[m':
 476             # Normal character attribute (all attributes off).
 477             self._attr = ATTR_NORMAL
 478             return True
 479         if seq == b'\x1b[7m':
 480             # Inverse character attribute.
 481             self._attr |= ATTR_REVERSE
 482             return True
 483         match = re.fullmatch(b'\x1b\\[3([0-9]+)m', seq)
 484         if match:
 485             # Set foreground color.
 486             color = int(match.group(1))
 487             if color in COLORS:
 488                 self._fgcolor = color
 489                 return True
 490             return False
 491         match = re.fullmatch(b'\x1b\\[4([0-9]+)m', seq)
 492         if match:
 493             # Set background color.
 494             color = int(match.group(1))
 495             if color in COLORS:
 496                 self._bgcolor = color
 497                 return True
 498             return False
 499         if seq == b'\x1b[?25l':
 500             # Hide cursor.
 501             return True
 502
 503         return False
 504
 505     def _cursor_down(self):
 506         """
 507         Move the screen cursor one line down. The screen is scrolled if the
 508         cursor points to the last line.
 509         """
 510
 511         if self._cur_y < ROWS - 1:
 512             self._cur_y += 1
 513         else:
 514             assert self._cur_y == ROWS - 1
 515
 516             # At the last line of the terminal, scroll up the screen.
 517             del self._screen[0]
 518             self._screen.append([TermChar() for x in range(COLUMNS)])
 519
 520             if self._root:
 521                 self._text.config(state=tkinter.NORMAL)
 522                 self._text.delete('1.0', '2.0')
 523                 self._text.insert(tkinter.END, "\n" + " " * COLUMNS)
 524                 self._text.config(state=tkinter.DISABLED)
 525
 526     def _erase_all(self):
 527         """Completely clear the terminal's screen."""
 528
 529         self._screen = [[TermChar() for x in range(COLUMNS)]
 530                         for y in range(ROWS)]
 531
 532         if self._root:
 533             self._text.config(state=tkinter.NORMAL)
 534             self._text.delete('1.0', tkinter.END)
 535             self._text.insert('1.0', "\n".join([" " * COLUMNS] * ROWS))
 536             self._text.config(state=tkinter.DISABLED)
 537
 538     def _insert_blanks(self, w):
 539         """
 540         Replace the specified number of characters on the current screen line
 541         with blanks.
 542         """
 543
 544         del self._screen[self._cur_y][-w:]
 545         pre = self._screen[self._cur_y][:self._cur_x]
 546         post = self._screen[self._cur_y][self._cur_x:]
 547         self._screen[self._cur_y] = pre + [TermChar() for x in range(w)] + post
 548
 549         if self._root:
 550             self._text.config(state=tkinter.NORMAL)
 551             self._text.delete('{}.end-{}c'.format(self._cur_y + 1, w),
 552                               '{}.end'.format(self._cur_y + 1))
 553             self._text.insert('{}.{}'.format(self._cur_y + 1, self._cur_x),
 554                               " " * w)
 555             self._text.config(state=tkinter.DISABLED)
 556
 557     def _print_char_at(self, y, x, char):
 558         """Output one character on the screen at the specified coordinates."""
 559
 560         # Record the character in the internal screen representation.
 561         self._screen[y][x] = char
 562
 563         if self._root:
 564             # Add the character to the terminal text widget.
 565             self._text.config(state=tkinter.NORMAL)
 566             pos = '{}.{}'.format(y + 1, x)
 567             self._text.delete(pos)
 568
 569             tag = 'tag_{}-{}'.format(char.get_tag_foreground(),
 570                                      char.get_tag_background())
 571             self._text.insert(pos, char.char, tag)
 572             self._text.config(state=tkinter.DISABLED)
 573
 574     def _print_char(self, char):
 575         """Output one character on the screen at the cursor position."""
 576
 577         self._print_char_at(self._cur_y, self._cur_x, char)
 578
 579         # Advance the cursor.
 580         self._cur_x += 1
 581         if self._cur_x == COLUMNS:
 582             self._cur_x = 0
 583             self._cursor_down()
 584
 585     def _tk_key(self, event):
 586         """Process a key pressed by the user."""
 587
 588         if len(event.char) != 0:
 589             if event.char == CODE_ENTER:
 590                 self._send_key(event.char, 'Enter')
 591             else:
 592                 self._send_key(event.char, event.char)
 593             return
 594
 595         # A special key was pressed.
 596         if event.keysym == 'F12':
 597             self._record_expected_screen()
 598             return
 599         if event.keysym == 'Prior':
 600             self._send_key(CODE_PAGE_UP, 'PageUp')
 601             return
 602         if event.keysym == 'Next':
 603             self._send_key(CODE_PAGE_DOWN, 'PageDown')
 604             return
 605         match = re.fullmatch('F([0-9]+)', event.keysym)
 606         if match:
 607             # F1 to F11.
 608             fnum = int(match.group(1))
 609             if fnum >= 1 and fnum <= len(CODE_FN):
 610                 self._send_key(CODE_FN[fnum - 1], event.keysym)
 611                 return
 612
 613         print("Unrecognized key {}.".format(event.keysym), file=sys.stderr)
 614
 615     def _get_screen_xml(self, screen):
 616         """
 617         Return an ElementTree.Element that represents the given screen.
 618         """
 619
 620         expect_e = ElementTree.Element('expect')
 621         data_e = ElementTree.SubElement(expect_e, 'data')
 622
 623         colors = {}
 624         new_key = 'a'
 625
 626         # Print content of the screen.
 627         for y in range(ROWS):
 628             line_e = ElementTree.SubElement(data_e, 'line')
 629             line_e.text = ''
 630
 631             attr = ''
 632
 633             for x in range(COLUMNS):
 634                 term_char = screen[y][x]
 635
 636                 line_e.text += term_char.char
 637
 638                 color = (term_char.attr, term_char.fgcolor, term_char.bgcolor)
 639                 if color == (ATTR_NORMAL, COLOR_DEFAULT, COLOR_DEFAULT):
 640                     key = ' '
 641                 elif color in colors:
 642                     key = colors[color]
 643                 else:
 644                     key = new_key
 645                     colors[color] = key
 646                     assert new_key != 'z'
 647                     new_key = chr(ord(new_key) + 1)
 648                 attr += key
 649
 650             # Record any non-default attributes/colors.
 651             if attr != ' ' * COLUMNS:
 652                 attr_e = ElementTree.SubElement(data_e, 'attr')
 653                 attr_e.text = attr
 654
 655         # Record used color schemes.
 656         if colors:
 657             scheme_e = ElementTree.SubElement(expect_e, 'scheme')
 658             for color, key in sorted(colors.items(), key=lambda x: x[1]):
 659                 attr, fgcolor, bgcolor = color
 660                 color_e = ElementTree.SubElement(scheme_e, 'color')
 661                 color_e.set('key', key)
 662
 663                 attr_str = attr_to_string(attr)
 664                 if attr_str:
 665                     color_e.set('attributes', attr_str)
 666
 667                 fgcolor_str = color_to_string(fgcolor)
 668                 if fgcolor_str:
 669                     color_e.set('foreground', fgcolor_str)
 670
 671                 bgcolor_str = color_to_string(bgcolor)
 672                 if bgcolor_str:
 673                     color_e.set('background', bgcolor_str)
 674
 675         return expect_e
 676
 677     def _record_expected_screen(self):
 678         """
 679         Record the current screen content as an expected screen in the test
 680         playbook that is being created.
 681         """
 682
 683         assert self._mode == self.MODE_RUN or self._mode == self.MODE_RECORD
 684
 685         if self._mode != self.MODE_RECORD:
 686             print("Recording is not enabled.", file=sys.stderr)
 687             return
 688
 689         expect_e = self._get_screen_xml(self._screen)
 690         self._test_e.append(expect_e)
 691         print("Recorded expected screen.")
 692
 693     # Method _indent_xml() is based on a code from
 694     # http://effbot.org/zone/element-lib.htm#prettyprint.
 695     def _indent_xml(self, elem, level=0):
 696         """
 697         Indent elements of a given ElementTree so it can be pretty-printed.
 698         """
 699
 700         i = '\n' + '\t' * level
 701         if len(elem):
 702             if not elem.text or not elem.text.strip():
 703                 elem.text = i + '\t'
 704             for e in elem:
 705                 self._indent_xml(e, level+1)
 706             if not e.tail or not e.tail.strip():
 707                 e.tail = i
 708         if not elem.tail or not elem.tail.strip():
 709             elem.tail = i
 710
 711     def output_test(self, filename):
 712         """
 713         Output a recorded playbook to a file with the given name. Returns True
 714         when the writing of the test data succeeded, False otherwise.
 715         """
 716
 717         assert self._mode == self.MODE_RECORD
 718
 719         # Pretty-format the XML tree.
 720         self._indent_xml(self._test_e)
 721
 722         # Output the test.
 723         tree = ElementTree.ElementTree(self._test_e)
 724         try:
 725             tree.write(filename, 'unicode', True)
 726         except Exception as e:
 727             print("Failed to write playbook file '{}': {}.".format(
 728                 filename, e), file=sys.stderr)
 729             return False
 730
 731         return True
 732
 733     class _TestFailure(Exception):
 734         """Exception reported when a test failed."""
 735         pass
 736
 737     def _playbook_key(self, cmd_e):
 738         """
 739         Parse a description of one key action and send the key to the terminal.
 740         Exception _TestFailure is raised if the description is malformed or
 741         incomplete, exception _TerminalConnectionException can be thrown when
 742         communication with the pseudo-terminal fails.
 743         """
 744
 745         assert self._mode == self.MODE_TEST
 746
 747         try:
 748             key = cmd_e.attrib['key']
 749         except KeyError:
 750             raise self._TestFailure("Element 'action' is missing required "
 751                                     "attribute 'key'")
 752
 753         # Handle simple characters.
 754         if len(key) == 1:
 755             self._send_key(key, key)
 756             return
 757
 758         # Handle special keys.
 759         if key == 'Enter':
 760             self._send_key(CODE_ENTER, key)
 761             return
 762         if key == 'PageUp':
 763             self._send_key(CODE_PAGE_UP, key)
 764             return
 765         if key == 'PageDown':
 766             self._send_key(CODE_PAGE_DOWN, key)
 767         match = re.fullmatch('F([0-9]+)', key)
 768         if match:
 769             # F1 to F11.
 770             fnum = int(match.group(1))
 771             if fnum >= 1 and fnum <= len(CODE_FN):
 772                 self._send_key(CODE_FN[fnum - 1], key)
 773                 return
 774
 775         raise self._TestFailure(
 776             "Element 'action' specifies unrecognized key '{}'".format(key))
 777
 778     def _parse_color_scheme(self, scheme_e):
 779         """
 780         Parse color scheme of one expected screen. Dictionary with
 781         {'key': (attr, fgcolor, bgcolor), ...} is returned on success,
 782         exception _TestFailure is raised if the description is malformed or
 783         incomplete.
 784         """
 785
 786         assert self._mode == self.MODE_TEST
 787
 788         colors = {}
 789         for color_e in scheme_e:
 790             try:
 791                 key = color_e.attrib['key']
 792             except KeyError:
 793                 raise self._TestFailure(
 794                     "Element 'color' is missing required attribute 'key'")
 795
 796             attr = ATTR_NORMAL
 797             if 'attributes' in color_e.attrib:
 798                 try:
 799                     attr = string_to_attr(color_e.attrib['attributes'])
 800                 except ValueError as e:
 801                     raise self._TestFailure(
 802                         "Value of attribute 'attributes' is invalid: "
 803                         "{}".format(e))
 804
 805             fgcolor = COLOR_DEFAULT
 806             if 'foreground' in color_e.attrib:
 807                 try:
 808                     fgcolor = string_to_color(color_e.attrib['foreground'])
 809                 except ValueError as e:
 810                     raise self._TestFailure(
 811                         "Value of attribute 'foreground' is invalid: "
 812                         "{}".format(e))
 813
 814             bgcolor = COLOR_DEFAULT
 815             if 'background' in color_e.attrib:
 816                 try:
 817                     bgcolor = string_to_color(color_e.attrib['background'])
 818                 except ValueError as e:
 819                     raise self._TestFailure(
 820                         "Value of attribute 'background' is invalid: "
 821                         "{}".format(e))
 822
 823             colors[key] = (attr, fgcolor, bgcolor)
 824
 825         return colors
 826
 827     def _parse_screen_data(self, data_e, colors):
 828         """
 829         Parse screen lines of one expected screen. Internal screen
 830         representation is returned on success, exception _TestFailure is raised
 831         if the description is malformed or incomplete.
 832         """
 833
 834         assert self._mode == self.MODE_TEST
 835
 836         NEW_LINE = 0
 837         NEW_LINE_OR_ATTR = 1
 838         state = NEW_LINE
 839         line = None
 840         expected_screen = []
 841
 842         for data_sub_e in data_e:
 843             # Do common processing for both states.
 844             if data_sub_e.tag == 'line':
 845                 # Append the previous line.
 846                 if line:
 847                     expected_screen.append(line)
 848                 # Parse the new line.
 849                 line = [TermChar(char) for char in data_sub_e.text]
 850                 state = NEW_LINE_OR_ATTR
 851
 852             if state == NEW_LINE and data_sub_e.tag != 'line':
 853                 raise self._TestFailure("Element '{}' is invalid, expected "
 854                                         "'line'".format(data_sub_e.tag))
 855
 856             elif state == NEW_LINE_OR_ATTR:
 857                 if data_sub_e.tag == 'attr':
 858                     if len(data_sub_e.text) != len(line):
 859                         raise self._TestFailure(
 860                             "Element 'attr' does not match the previous line, "
 861                             "expected '{}' attribute characters but got "
 862                             "'{}'".format(len(line), len(data_sub_e.text)))
 863
 864                     for i, key in enumerate(data_sub_e.text):
 865                         if key == ' ':
 866                             continue
 867
 868                         try:
 869                             attr, fgcolor, bgcolor = colors[key]
 870                         except KeyError:
 871                             raise self._TestFailure("Color attribute '{}' is "
 872                                                     "not defined".format(key))
 873                         line[i].attr = attr
 874                         line[i].fgcolor = fgcolor
 875                         line[i].bgcolor = bgcolor
 876
 877                     state = NEW_LINE
 878
 879                 elif data_sub_e.tag != 'line':
 880                     raise self._TestFailure(
 881                         "Element '{}' is invalid, expected 'line' or "
 882                         "'attr'".format(data_sub_e.tag))
 883
 884         # Append the final line.
 885         if line:
 886             expected_screen.append(line)
 887
 888         return expected_screen
 889
 890     def _parse_expected_screen(self, expect_e):
 891         """
 892         Parse a description of one expected screen. Internal screen
 893         representation is returned on success, exception _TestFailure is raised
 894         if the description is malformed or incomplete.
 895         """
 896
 897         assert self._mode == self.MODE_TEST
 898
 899         data_e = None
 900         scheme_e = None
 901         for sub_e in expect_e:
 902             if sub_e.tag == 'data':
 903                 if data_e:
 904                     raise self._TestFailure("Element 'expect' contains "
 905                                             "multiple 'data' sub-elements")
 906                 data_e = sub_e
 907             elif sub_e.tag == 'scheme':
 908                 if scheme_e:
 909                     raise self._TestFailure("Element 'expect' contains "
 910                                             "multiple 'scheme' sub-elements")
 911                 scheme_e = sub_e
 912
 913         if not data_e:
 914             raise self._TestFailure(
 915                 "Element 'expect' is missing required sub-element 'data'")
 916
 917         # Parse the color scheme.
 918         if scheme_e:
 919             colors = self._parse_color_scheme(scheme_e)
 920         else:
 921             colors = {}
 922
 923         # Parse the screen data.
 924         return self._parse_screen_data(data_e, colors)
 925
 926     def _report_failed_expectation(self, expected_screen):
 927         """
 928         Report that the expected screen state has not been reached. The output
 929         consists of the expected screen, the current screen content, followed
 930         by differences between the two screens.
 931         """
 932
 933         assert self._mode == self.MODE_TEST
 934
 935         # Print the expected screen. The output is not verbatim as it was
 936         # specified in the input file, but instead the screen is printed in the
 937         # same way the current screen gets output. This allows to properly show
 938         # differences between the two screens.
 939
 940         expected_screen_e = self._get_screen_xml(expected_screen)
 941         self._indent_xml(expected_screen_e)
 942         expected_screen_str = ElementTree.tostring(
 943             expected_screen_e, 'unicode')
 944         print("Expected (normalized) screen:", file=sys.stderr)
 945         print(expected_screen_str, file=sys.stderr)
 946
 947         # Print the current screen.
 948         current_screen_e = self._get_screen_xml(self._screen)
 949         self._indent_xml(current_screen_e)
 950         current_screen_str = ElementTree.tostring(current_screen_e, 'unicode')
 951         print("Current screen:", file=sys.stderr)
 952         print(current_screen_str, file=sys.stderr)
 953
 954         # Print the delta.
 955         print("Differences:", file=sys.stderr)
 956         sys.stderr.writelines(difflib.unified_diff(
 957             expected_screen_str.splitlines(keepends=True),
 958             current_screen_str.splitlines(keepends=True),
 959             fromfile="Expected screen", tofile="Current screen"))
 960
 961     def _execute_playbook(self, test_e):
 962         """
 963         Run the main loop and execute the given test playbook. Normal return
 964         from the method indicates that the test succeeded. Exception
 965         _TestFailure is raised when the test fails and exception
 966         _TerminalConnectionException can be thrown when communication with the
 967         pseudo-terminal fails.
 968         """
 969
 970         assert self._mode == self.MODE_TEST
 971
 972         cmd_iter = iter(test_e)
 973
 974         # Start the main loop.
 975         with selectors.DefaultSelector() as sel:
 976             sel.register(self._fd, selectors.EVENT_READ)
 977
 978             expected_screen = None
 979             more_commands = True
 980             while True:
 981                 # Process any actions and find an expected screen.
 982                 while not expected_screen and more_commands:
 983                     try:
 984                         cmd_e = next(cmd_iter)
 985                         if cmd_e.tag == 'action':
 986                             self._playbook_key(cmd_e)
 987                         elif cmd_e.tag == 'expect':
 988                             expected_screen = self._parse_expected_screen(
 989                                 cmd_e)
 990                             # Stop processing more commands for now and wait
 991                             # for the expected screen to appear.
 992                             break
 993                         else:
 994                             raise self._TestFailure(
 995                                 "Element '{}' is invalid, expected 'action' "
 996                                 "or 'expect'".format(cmd_e.tag))
 997                     except StopIteration:
 998                         # No more commands.
 999                         more_commands = False
1000
1001                 # Wait for the expected screen.
1002                 events = sel.select(CHILD_TIMEOUT)
1003                 if not events:
1004                     if expected_screen:
1005                         self._report_failed_expectation(expected_screen)
1006                     raise self._TestFailure(
1007                         "Timeout reached. No event received in the last {} "
1008                         "second(s)".format(CHILD_TIMEOUT))
1009
1010                 # Expect only an event on self._fd.
1011                 assert len(events) == 1
1012                 event = events[0]
1013                 key, _mask = event
1014                 assert key.fd == self._fd
1015
1016                 closed = self._pty_callback()
1017                 if closed:
1018                     if more_commands:
1019                         raise self._TestFailure(
1020                             "Connection to the terminal was closed but the "
1021                             "playbook contains more commands")
1022                     break
1023
1024                 # Check if the expected screen is present.
1025                 if self._screen == expected_screen:
1026                     expected_screen = None
1027
1028     def execute_test(self, filename):
1029         """
1030         Load test data from a given file, start the program under the test and
1031         execute the test playbook. Returns True when the test succeeded, False
1032         otherwise.
1033         """
1034
1035         assert self._mode == self.MODE_TEST
1036
1037         # Read the test data.
1038         try:
1039             tree = ElementTree.ElementTree(file=filename)
1040         except Exception as e:
1041             print("Failed to read playbook file '{}': {}.".format(filename, e),
1042                   file=sys.stderr)
1043             return False
1044
1045         # Read what program to execute.
1046         test_e = tree.getroot()
1047         if test_e.tag != 'test':
1048             print("Root element '{}' is invalid, expected 'test'.".format(
1049                   test_e.tag), file=sys.stderr)
1050             return False
1051         try:
1052             self._program = test_e.attrib['program']
1053         except KeyError:
1054             print("Element 'test' is missing required attribute 'program'.",
1055                   file=sys.stderr)
1056             return False
1057
1058         # Start the specified program.
1059         if not self._start_program():
1060             return False
1061
1062         # Execute the test playbook.
1063         res = True
1064         try:
1065             self._execute_playbook(tree.getroot())
1066         except (self._TerminalConnectionException, self._TestFailure) as e:
1067             print("{}.".format(e), file=sys.stderr)
1068             res = False
1069         finally:
1070             # Finalize the run of the child program.
1071             if not self._finalize_program():
1072                 res = False
1073
1074         # Return whether the test passed.
1075         return res
1076
1077
1078 def main():
1079     """
1080     Parse command line arguments and execute the operation that the user
1081     selected. Returns 0 if the operation was successful and a non-zero value
1082     otherwise.
1083     """
1084
1085     # Parse command line arguments.
1086     parser = argparse.ArgumentParser()
1087     parser.add_argument(
1088         '-t', '--terminfo', metavar='PATH', help="path to terminfo directory")
1089
1090     subparsers = parser.add_subparsers(dest='program')
1091     subparsers.required = True
1092
1093     program_parser = argparse.ArgumentParser(add_help=False)
1094     program_parser.add_argument('program', help="executable to run")
1095
1096     # Create the parser for the 'run' command.
1097     parser_run = subparsers.add_parser(
1098         'run', parents=[program_parser], help="run a program")
1099     parser_run.set_defaults(mode=Term.MODE_RUN)
1100
1101     # Create the parser for the 'record' command.
1102     parser_record = subparsers.add_parser(
1103         'record', parents=[program_parser], help="record a test")
1104     parser_record.set_defaults(mode=Term.MODE_RECORD)
1105     parser_record.add_argument(
1106         '-o', '--playbook', metavar='FILE', required=True,
1107         help="output playbook file")
1108
1109     # Create the parser for the 'test' command.
1110     parser_test = subparsers.add_parser('test', help="perform a test")
1111     parser_test.set_defaults(program=None)
1112     parser_test.set_defaults(mode=Term.MODE_TEST)
1113     parser_test.add_argument('playbook', help="input playbook file")
1114
1115     args = parser.parse_args()
1116
1117     tk_root = None
1118     if args.mode in (Term.MODE_RUN, Term.MODE_RECORD):
1119         # Start the terminal GUI.
1120         try:
1121             tk_root = tkinter.Tk()
1122         except tkinter.TclError as e:
1123             print("Failed to initialize GUI: {}.".format(e), file=sys.stderr)
1124             return 1
1125
1126     term = Term(tk_root, args.program, args.mode, args.terminfo)
1127     if tk_root:
1128         # Start the GUI main loop.
1129         term.run_gui_mainloop()
1130     else:
1131         # Execute and check the playbook, without running GUI.
1132         ok = term.execute_test(args.playbook)
1133         if ok:
1134             msg = "succeeded"
1135             res = 0
1136         else:
1137             msg = "failed"
1138             res = 1
1139
1140         print("Checking of playbook '{}' {}.".format(args.playbook, msg))
1141         return res
1142
1143     if args.mode == Term.MODE_RECORD:
1144         # Get the recorded test data and write them to a file.
1145         if not term.output_test(args.playbook):
1146             return 1
1147
1148     return 0
1149
1150
1151 if __name__ == '__main__':
1152     sys.exit(main())
1153
1154 # vim: set tabstop=4 shiftwidth=4 textwidth=79 expandtab :