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 xml.etree.ElementTree as ElementTree
  32
  33 # The module relies on selectors.select() to automatically retry the operation
  34 # with a recomputed timeout when it gets interrupted by a signal. This
  35 # behaviour was introduced in Python 3.5.
  36 if sys.hexversion < 0x03050000:
  37     print("This program requires at least Python 3.5.", file=sys.stderr)
  38     sys.exit(1)
  39
  40 ROWS = 24
  41 COLUMNS = 80
  42 CHILD_TIMEOUT = 5
  43
  44 ATTR_NORMAL = 0
  45 ATTR_REVERSE = 1
  46
  47 COLOR_BLACK = 0
  48 COLOR_RED = 1
  49 COLOR_GREEN = 2
  50 COLOR_YELLOW = 3
  51 COLOR_BLUE = 4
  52 COLOR_MAGENTA = 5
  53 COLOR_CYAN = 6
  54 COLOR_WHITE = 7
  55 COLOR_DEFAULT = 9
  56 COLOR_REGISTER = ((COLOR_BLACK, 'black'), (COLOR_RED, 'red'),
  57                   (COLOR_GREEN, 'green'), (COLOR_YELLOW, 'yellow'),
  58                   (COLOR_BLUE, 'blue'), (COLOR_MAGENTA, 'magenta'),
  59                   (COLOR_CYAN, 'cyan'), (COLOR_WHITE, 'white'),
  60                   (COLOR_DEFAULT, 'default'))
  61 COLOR_TO_STRING_MAP = {id_: name for id_, name in COLOR_REGISTER}
  62 STRING_TO_COLOR_MAP = {name: id_ for id_, name in COLOR_REGISTER}
  63 COLORS = {id_ for id_, _ in COLOR_REGISTER}
  64 REAL_COLOR_NAMES = tuple(
  65     name for id_, name in COLOR_REGISTER if id_ != COLOR_DEFAULT)
  66
  67 COLOR_DEFAULT_FOREGROUND = COLOR_BLACK
  68 COLOR_DEFAULT_BACKGROUND = COLOR_WHITE
  69
  70 CODE_ENTER = '\x0d'
  71 CODE_FN = ('\x1bOP', '\x1bOQ', '\x1bOR', '\x1bOS', '\x1b[15~', '\x1b[17~',
  72            '\x1b[18~', '\x1b[19~', '\x1b[20~', '\x1b[21~', '\x1b[23~')
  73 CODE_PAGE_UP = '\x1b[5~'
  74 CODE_PAGE_DOWN = '\x1b[6~'
  75
  76
  77 def attr_to_string(attr):
  78     """Get string representation of given attributes."""
  79
  80     res = []
  81     if (attr & ATTR_REVERSE) != 0:
  82         res.append('reverse')
  83     return '|'.join(res)
  84
  85
  86 def string_to_attr(string):
  87     """
  88     Convert a string to attributes. Exception ValueError is raised if some
  89     attribute is invalid.
  90     """
  91
  92     res = ATTR_NORMAL
  93     for attr in string.split('|'):
  94         if attr == 'normal':
  95             pass
  96         elif attr == 'reverse':
  97             res |= ATTR_REVERSE
  98         else:
  99             raise ValueError("Unrecognized attribute '{}'".format(attr))
 100     return res
 101
 102
 103 def color_to_string(color):
 104     """Get string representation of a given color."""
 105
 106     return COLOR_TO_STRING_MAP[color]
 107
 108
 109 def string_to_color(string):
 110     """
 111     Convert a string to a color. Exception ValueError is raised if the color
 112     name is not recognized.
 113     """
 114
 115     try:
 116         return STRING_TO_COLOR_MAP[string]
 117     except KeyError:
 118         raise ValueError("Unrecognized color '{}'".format(string))
 119
 120
 121 class TermChar:
 122     """On-screen character."""
 123
 124     def __init__(self, char=" ", attr=ATTR_NORMAL, fgcolor=COLOR_DEFAULT,
 125                  bgcolor=COLOR_DEFAULT):
 126         self.char = char
 127         self.attr = attr
 128         self.fgcolor = fgcolor
 129         self.bgcolor = bgcolor
 130
 131     def __eq__(self, other):
 132         return self.__dict__ == other.__dict__
 133
 134     def _get_translated_fgcolor(self):
 135         """
 136         Return the foreground color. If the current color is COLOR_DEFAULT then
 137         COLOR_DEFAULT_FOREGROUND is returned.
 138         """
 139
 140         if self.fgcolor == COLOR_DEFAULT:
 141             return COLOR_DEFAULT_FOREGROUND
 142         return self.fgcolor
 143
 144     def _get_translated_bgcolor(self):
 145         """
 146         Return the background color. If the current color is COLOR_DEFAULT then
 147         COLOR_DEFAULT_BACKGROUND is returned.
 148         """
 149
 150         if self.bgcolor == COLOR_DEFAULT:
 151             return COLOR_DEFAULT_BACKGROUND
 152         return self.bgcolor
 153
 154     def get_tag_foreground(self):
 155         """
 156         Return a name of the final foreground color that should be used to
 157         display the character on the screen.
 158         """
 159
 160         if self.attr & ATTR_REVERSE:
 161             color = self._get_translated_bgcolor()
 162         else:
 163             color = self._get_translated_fgcolor()
 164         return color_to_string(color)
 165
 166     def get_tag_background(self):
 167         """
 168         Return a name of the final background color that should be used to
 169         display the character on the screen.
 170         """
 171
 172         if self.attr & ATTR_REVERSE:
 173             color = self._get_translated_fgcolor()
 174         else:
 175             color = self._get_translated_bgcolor()
 176         return color_to_string(color)
 177
 178
 179 class Term:
 180     """Termex terminal emulator."""
 181
 182     MODE_RUN = 0
 183     MODE_RECORD = 1
 184     MODE_TEST = 2
 185
 186     class _TerminalConnectionException(Exception):
 187         """
 188         Exception reported when communication with the pseudo-terminal fails.
 189         """
 190         pass
 191
 192     def __init__(self, root, program, mode, terminfo=None):
 193         self._root = root
 194         self._program = program
 195         self._mode = mode
 196         self._terminfo = terminfo
 197
 198         # Test mode obtains the program name from the playbook.
 199         if self._mode == self.MODE_TEST:
 200             assert self._program is None
 201
 202         self._child_pid = None
 203         self._fd = None
 204
 205         self._screen = None
 206         self._cur_y = 0
 207         self._cur_x = 0
 208         self._attr = ATTR_NORMAL
 209         self._fgcolor = COLOR_DEFAULT
 210         self._bgcolor = COLOR_DEFAULT
 211         self._charbuf = b''
 212         self._alt_charset_mode = False
 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         # Handle characters from the alternate set.
 418         if self._alt_charset_mode:
 419             uchars = {
 420                 b'q': '\N{BOX DRAWINGS LIGHT HORIZONTAL}',
 421                 b'x': '\N{BOX DRAWINGS LIGHT VERTICAL}',
 422                 b'm': '\N{BOX DRAWINGS LIGHT UP AND RIGHT}',
 423                 b'j': '\N{BOX DRAWINGS LIGHT UP AND LEFT}',
 424                 b'l': '\N{BOX DRAWINGS LIGHT DOWN AND RIGHT}',
 425                 b'k': '\N{BOX DRAWINGS LIGHT DOWN AND LEFT}',
 426                 b'v': '\N{BOX DRAWINGS LIGHT UP AND HORIZONTAL}',
 427                 b't': '\N{BOX DRAWINGS LIGHT VERTICAL AND RIGHT}',
 428                 b'u': '\N{BOX DRAWINGS LIGHT VERTICAL AND LEFT}',
 429                 b'w': '\N{BOX DRAWINGS LIGHT DOWN AND HORIZONTAL}',
 430                 b'.': '\N{DOWNWARDS ARROW}',
 431                 b',': '\N{LEFTWARDS ARROW}',
 432                 b'+': '\N{RIGHTWARDS ARROW}',
 433                 b'-': '\N{UPWARDS ARROW}',
 434                 b'~': '\N{MIDDLE DOT}'}
 435
 436             if seq in uchars:
 437                 self._print_char(TermChar(uchars[seq], self._attr,
 438                                           self._fgcolor, self._bgcolor))
 439                 return True
 440
 441         # Handle normal characters.
 442         if re.fullmatch(b'[^\x01-\x1f]+', seq):
 443             try:
 444                 uchar = seq.decode('utf-8')
 445                 self._print_char(TermChar(uchar, self._attr, self._fgcolor,
 446                                           self._bgcolor))
 447                 return True
 448             except UnicodeError:
 449                 # Continue on the assumption that it is not yet a complete
 450                 # character. This assumption is wrong if the received text is
 451                 # actually malformed.
 452                 return False
 453
 454         if seq == b'\x07':
 455             # Bell.
 456             if self._root:
 457                 self._root.bell()
 458             return True
 459         if seq == b'\x08':
 460             # Backspace non-destructively.
 461             self._cur_x -= 1
 462             return True
 463         if seq == b'\x0d':
 464             # Go to beginning of line.
 465             self._cur_x = 0
 466             return True
 467         if seq == b'\x0a':
 468             # Move cursor down one line.
 469             self._cursor_down()
 470             return True
 471
 472         # Controls beginning with ESC.
 473         if seq == b'\x1b(0':
 474             # Start alternate character set.
 475             self._alt_charset_mode = True
 476             return True
 477         if seq == b'\x1b(B':
 478             # End alternate character set.
 479             self._alt_charset_mode = False
 480             return True
 481
 482         # Control sequences.
 483         match = re.fullmatch(b'\x1b\\[([0-9]+)@', seq)
 484         if match:
 485             # Insert blank characters.
 486             self._insert_blanks(int(match.group(1)))
 487             return True
 488         if seq == b'\x1b[H':
 489             # Set cursor position to the default (top left).
 490             self._cur_y = 0
 491             self._cur_x = 0
 492             return True
 493         match = re.fullmatch(b'\x1b\\[([0-9]+);([0-9]+)H', seq)
 494         if match:
 495             # Set cursor position to (y,x).
 496             self._cur_y = int(match.group(1))
 497             self._cur_x = int(match.group(2))
 498             return True
 499         if self._charbuf == b'\x1b[K':
 500             # Erase in line to right.
 501             for x in range(self._cur_x, COLUMNS):
 502                 self._print_char_at(self._cur_y, x, TermChar())
 503             return True
 504         if seq == b'\x1b[2J':
 505             # Erase display completely.
 506             self._erase_all()
 507             return True
 508         if seq == b'\x1b[m':
 509             # Normal character attribute (all attributes off).
 510             self._attr = ATTR_NORMAL
 511             return True
 512         if seq == b'\x1b[7m':
 513             # Inverse character attribute.
 514             self._attr |= ATTR_REVERSE
 515             return True
 516         match = re.fullmatch(b'\x1b\\[3([0-9]+)m', seq)
 517         if match:
 518             # Set foreground color.
 519             color = int(match.group(1))
 520             if color in COLORS:
 521                 self._fgcolor = color
 522                 return True
 523             return False
 524         match = re.fullmatch(b'\x1b\\[4([0-9]+)m', seq)
 525         if match:
 526             # Set background color.
 527             color = int(match.group(1))
 528             if color in COLORS:
 529                 self._bgcolor = color
 530                 return True
 531             return False
 532         if seq == b'\x1b[?25l':
 533             # Hide cursor.
 534             return True
 535
 536         return False
 537
 538     def _cursor_down(self):
 539         """
 540         Move the screen cursor one line down. The screen is scrolled if the
 541         cursor points to the last line.
 542         """
 543
 544         if self._cur_y < ROWS - 1:
 545             self._cur_y += 1
 546         else:
 547             assert self._cur_y == ROWS - 1
 548
 549             # At the last line of the terminal, scroll up the screen.
 550             del self._screen[0]
 551             self._screen.append([TermChar() for x in range(COLUMNS)])
 552
 553             if self._root:
 554                 self._text.config(state=tkinter.NORMAL)
 555                 self._text.delete('1.0', '2.0')
 556                 self._text.insert(tkinter.END, "\n" + " " * COLUMNS)
 557                 self._text.config(state=tkinter.DISABLED)
 558
 559     def _erase_all(self):
 560         """Completely clear the terminal's screen."""
 561
 562         self._screen = [[TermChar() for x in range(COLUMNS)]
 563                         for y in range(ROWS)]
 564
 565         if self._root:
 566             self._text.config(state=tkinter.NORMAL)
 567             self._text.delete('1.0', tkinter.END)
 568             self._text.insert('1.0', "\n".join([" " * COLUMNS] * ROWS))
 569             self._text.config(state=tkinter.DISABLED)
 570
 571     def _insert_blanks(self, w):
 572         """
 573         Replace the specified number of characters on the current screen line
 574         with blanks.
 575         """
 576
 577         del self._screen[self._cur_y][-w:]
 578         pre = self._screen[self._cur_y][:self._cur_x]
 579         post = self._screen[self._cur_y][self._cur_x:]
 580         self._screen[self._cur_y] = pre + [TermChar() for x in range(w)] + post
 581
 582         if self._root:
 583             self._text.config(state=tkinter.NORMAL)
 584             self._text.delete('{}.end-{}c'.format(self._cur_y + 1, w),
 585                               '{}.end'.format(self._cur_y + 1))
 586             self._text.insert('{}.{}'.format(self._cur_y + 1, self._cur_x),
 587                               " " * w)
 588             self._text.config(state=tkinter.DISABLED)
 589
 590     def _print_char_at(self, y, x, char):
 591         """Output one character on the screen at the specified coordinates."""
 592
 593         # Record the character in the internal screen representation.
 594         self._screen[y][x] = char
 595
 596         if self._root:
 597             # Add the character to the terminal text widget.
 598             self._text.config(state=tkinter.NORMAL)
 599             pos = '{}.{}'.format(y + 1, x)
 600             self._text.delete(pos)
 601
 602             tag = 'tag_{}-{}'.format(char.get_tag_foreground(),
 603                                      char.get_tag_background())
 604             self._text.insert(pos, char.char, tag)
 605             self._text.config(state=tkinter.DISABLED)
 606
 607     def _print_char(self, char):
 608         """Output one character on the screen at the cursor position."""
 609
 610         self._print_char_at(self._cur_y, self._cur_x, char)
 611
 612         # Advance the cursor.
 613         self._cur_x += 1
 614         if self._cur_x == COLUMNS:
 615             self._cur_x = 0
 616             self._cursor_down()
 617
 618     def _tk_key(self, event):
 619         """Process a key pressed by the user."""
 620
 621         if len(event.char) != 0:
 622             if event.char == CODE_ENTER:
 623                 self._send_key(event.char, 'Enter')
 624             else:
 625                 self._send_key(event.char, event.char)
 626             return
 627
 628         # A special key was pressed.
 629         if event.keysym == 'F12':
 630             self._record_expected_screen()
 631             return
 632         if event.keysym == 'Prior':
 633             self._send_key(CODE_PAGE_UP, 'PageUp')
 634             return
 635         if event.keysym == 'Next':
 636             self._send_key(CODE_PAGE_DOWN, 'PageDown')
 637             return
 638         match = re.fullmatch('F([0-9]+)', event.keysym)
 639         if match:
 640             # F1 to F11.
 641             fnum = int(match.group(1))
 642             if fnum >= 1 and fnum <= len(CODE_FN):
 643                 self._send_key(CODE_FN[fnum - 1], event.keysym)
 644                 return
 645
 646         print("Unrecognized key {}.".format(event.keysym), file=sys.stderr)
 647
 648     def _get_screen_xml(self, screen):
 649         """
 650         Return an ElementTree.Element that represents the given screen.
 651         """
 652
 653         expect_e = ElementTree.Element('expect')
 654         data_e = ElementTree.SubElement(expect_e, 'data')
 655
 656         colors = {}
 657         new_key = 'a'
 658
 659         # Print content of the screen.
 660         for y in range(ROWS):
 661             line_e = ElementTree.SubElement(data_e, 'line')
 662             line_e.text = ''
 663
 664             attr = ''
 665
 666             for x in range(COLUMNS):
 667                 term_char = screen[y][x]
 668
 669                 line_e.text += term_char.char
 670
 671                 color = (term_char.attr, term_char.fgcolor, term_char.bgcolor)
 672                 if color == (ATTR_NORMAL, COLOR_DEFAULT, COLOR_DEFAULT):
 673                     key = ' '
 674                 elif color in colors:
 675                     key = colors[color]
 676                 else:
 677                     key = new_key
 678                     colors[color] = key
 679                     assert new_key != 'z'
 680                     new_key = chr(ord(new_key) + 1)
 681                 attr += key
 682
 683             # Record any non-default attributes/colors.
 684             if attr != ' ' * COLUMNS:
 685                 attr_e = ElementTree.SubElement(data_e, 'attr')
 686                 attr_e.text = attr
 687
 688         # Record used color schemes.
 689         if colors:
 690             scheme_e = ElementTree.SubElement(expect_e, 'scheme')
 691             for color, key in sorted(colors.items(), key=lambda x: x[1]):
 692                 attr, fgcolor, bgcolor = color
 693                 color_e = ElementTree.SubElement(scheme_e, 'color')
 694                 color_e.set('key', key)
 695
 696                 attr_str = attr_to_string(attr)
 697                 if attr_str:
 698                     color_e.set('attributes', attr_str)
 699
 700                 fgcolor_str = color_to_string(fgcolor)
 701                 if fgcolor_str:
 702                     color_e.set('foreground', fgcolor_str)
 703
 704                 bgcolor_str = color_to_string(bgcolor)
 705                 if bgcolor_str:
 706                     color_e.set('background', bgcolor_str)
 707
 708         return expect_e
 709
 710     def _record_expected_screen(self):
 711         """
 712         Record the current screen content as an expected screen in the test
 713         playbook that is being created.
 714         """
 715
 716         assert self._mode == self.MODE_RUN or self._mode == self.MODE_RECORD
 717
 718         if self._mode != self.MODE_RECORD:
 719             print("Recording is not enabled.", file=sys.stderr)
 720             return
 721
 722         expect_e = self._get_screen_xml(self._screen)
 723         self._test_e.append(expect_e)
 724         print("Recorded expected screen.")
 725
 726     # Method _indent_xml() is based on a code from
 727     # http://effbot.org/zone/element-lib.htm#prettyprint.
 728     def _indent_xml(self, elem, level=0):
 729         """
 730         Indent elements of a given ElementTree so it can be pretty-printed.
 731         """
 732
 733         i = '\n' + '\t' * level
 734         if len(elem):
 735             if not elem.text or not elem.text.strip():
 736                 elem.text = i + '\t'
 737             for e in elem:
 738                 self._indent_xml(e, level+1)
 739             if not e.tail or not e.tail.strip():
 740                 e.tail = i
 741         if not elem.tail or not elem.tail.strip():
 742             elem.tail = i
 743
 744     def output_test(self, filename):
 745         """
 746         Output a recorded playbook to a file with the given name. Returns True
 747         when the writing of the test data succeeded, False otherwise.
 748         """
 749
 750         assert self._mode == self.MODE_RECORD
 751
 752         # Pretty-format the XML tree.
 753         self._indent_xml(self._test_e)
 754
 755         # Output the test.
 756         tree = ElementTree.ElementTree(self._test_e)
 757         try:
 758             tree.write(filename, 'unicode', True)
 759         except Exception as e:
 760             print("Failed to write playbook file '{}': {}.".format(
 761                 filename, e), file=sys.stderr)
 762             return False
 763
 764         return True
 765
 766     class _TestFailure(Exception):
 767         """Exception reported when a test failed."""
 768         pass
 769
 770     def _playbook_key(self, cmd_e):
 771         """
 772         Parse a description of one key action and send the key to the terminal.
 773         Exception _TestFailure is raised if the description is malformed or
 774         incomplete, exception _TerminalConnectionException can be thrown when
 775         communication with the pseudo-terminal fails.
 776         """
 777
 778         assert self._mode == self.MODE_TEST
 779
 780         try:
 781             key = cmd_e.attrib['key']
 782         except KeyError:
 783             raise self._TestFailure("Element 'action' is missing required "
 784                                     "attribute 'key'")
 785
 786         # Handle simple characters.
 787         if len(key) == 1:
 788             self._send_key(key, key)
 789             return
 790
 791         # Handle special keys.
 792         if key == 'Enter':
 793             self._send_key(CODE_ENTER, key)
 794             return
 795         if key == 'PageUp':
 796             self._send_key(CODE_PAGE_UP, key)
 797             return
 798         if key == 'PageDown':
 799             self._send_key(CODE_PAGE_DOWN, key)
 800         match = re.fullmatch('F([0-9]+)', key)
 801         if match:
 802             # F1 to F11.
 803             fnum = int(match.group(1))
 804             if fnum >= 1 and fnum <= len(CODE_FN):
 805                 self._send_key(CODE_FN[fnum - 1], key)
 806                 return
 807
 808         raise self._TestFailure(
 809             "Element 'action' specifies unrecognized key '{}'".format(key))
 810
 811     def _parse_color_scheme(self, scheme_e):
 812         """
 813         Parse color scheme of one expected screen. Dictionary with
 814         {'key': (attr, fgcolor, bgcolor), ...} is returned on success,
 815         exception _TestFailure is raised if the description is malformed or
 816         incomplete.
 817         """
 818
 819         assert self._mode == self.MODE_TEST
 820
 821         colors = {}
 822         for color_e in scheme_e:
 823             try:
 824                 key = color_e.attrib['key']
 825             except KeyError:
 826                 raise self._TestFailure(
 827                     "Element 'color' is missing required attribute 'key'")
 828
 829             attr = ATTR_NORMAL
 830             if 'attributes' in color_e.attrib:
 831                 try:
 832                     attr = string_to_attr(color_e.attrib['attributes'])
 833                 except ValueError as e:
 834                     raise self._TestFailure(
 835                         "Value of attribute 'attributes' is invalid: "
 836                         "{}".format(e))
 837
 838             fgcolor = COLOR_DEFAULT
 839             if 'foreground' in color_e.attrib:
 840                 try:
 841                     fgcolor = string_to_color(color_e.attrib['foreground'])
 842                 except ValueError as e:
 843                     raise self._TestFailure(
 844                         "Value of attribute 'foreground' is invalid: "
 845                         "{}".format(e))
 846
 847             bgcolor = COLOR_DEFAULT
 848             if 'background' in color_e.attrib:
 849                 try:
 850                     bgcolor = string_to_color(color_e.attrib['background'])
 851                 except ValueError as e:
 852                     raise self._TestFailure(
 853                         "Value of attribute 'background' is invalid: "
 854                         "{}".format(e))
 855
 856             colors[key] = (attr, fgcolor, bgcolor)
 857
 858         return colors
 859
 860     def _parse_screen_data(self, data_e, colors):
 861         """
 862         Parse screen lines of one expected screen. Internal screen
 863         representation is returned on success, exception _TestFailure is raised
 864         if the description is malformed or incomplete.
 865         """
 866
 867         assert self._mode == self.MODE_TEST
 868
 869         NEW_LINE = 0
 870         NEW_LINE_OR_ATTR = 1
 871         state = NEW_LINE
 872         line = None
 873         expected_screen = []
 874
 875         for data_sub_e in data_e:
 876             # Do common processing for both states.
 877             if data_sub_e.tag == 'line':
 878                 # Append the previous line.
 879                 if line:
 880                     expected_screen.append(line)
 881                 # Parse the new line.
 882                 line = [TermChar(char) for char in data_sub_e.text]
 883                 state = NEW_LINE_OR_ATTR
 884
 885             if state == NEW_LINE and data_sub_e.tag != 'line':
 886                 raise self._TestFailure("Element '{}' is invalid, expected "
 887                                         "'line'".format(data_sub_e.tag))
 888
 889             elif state == NEW_LINE_OR_ATTR:
 890                 if data_sub_e.tag == 'attr':
 891                     if len(data_sub_e.text) != len(line):
 892                         raise self._TestFailure(
 893                             "Element 'attr' does not match the previous line, "
 894                             "expected '{}' attribute characters but got "
 895                             "'{}'".format(len(line), len(data_sub_e.text)))
 896
 897                     for i, key in enumerate(data_sub_e.text):
 898                         if key == ' ':
 899                             continue
 900
 901                         try:
 902                             attr, fgcolor, bgcolor = colors[key]
 903                         except KeyError:
 904                             raise self._TestFailure("Color attribute '{}' is "
 905                                                     "not defined".format(key))
 906                         line[i].attr = attr
 907                         line[i].fgcolor = fgcolor
 908                         line[i].bgcolor = bgcolor
 909
 910                     state = NEW_LINE
 911
 912                 elif data_sub_e.tag != 'line':
 913                     raise self._TestFailure(
 914                         "Element '{}' is invalid, expected 'line' or "
 915                         "'attr'".format(data_sub_e.tag))
 916
 917         # Append the final line.
 918         if line:
 919             expected_screen.append(line)
 920
 921         return expected_screen
 922
 923     def _parse_expected_screen(self, expect_e):
 924         """
 925         Parse a description of one expected screen. Internal screen
 926         representation is returned on success, exception _TestFailure is raised
 927         if the description is malformed or incomplete.
 928         """
 929
 930         assert self._mode == self.MODE_TEST
 931
 932         data_e = None
 933         scheme_e = None
 934         for sub_e in expect_e:
 935             if sub_e.tag == 'data':
 936                 if data_e:
 937                     raise self._TestFailure("Element 'expect' contains "
 938                                             "multiple 'data' sub-elements")
 939                 data_e = sub_e
 940             elif sub_e.tag == 'scheme':
 941                 if scheme_e:
 942                     raise self._TestFailure("Element 'expect' contains "
 943                                             "multiple 'scheme' sub-elements")
 944                 scheme_e = sub_e
 945
 946         if not data_e:
 947             raise self._TestFailure(
 948                 "Element 'expect' is missing required sub-element 'data'")
 949
 950         # Parse the color scheme.
 951         if scheme_e:
 952             colors = self._parse_color_scheme(scheme_e)
 953         else:
 954             colors = {}
 955
 956         # Parse the screen data.
 957         return self._parse_screen_data(data_e, colors)
 958
 959     def _report_failed_expectation(self, expected_screen):
 960         """
 961         Report that the expected screen state has not been reached. The output
 962         consists of the expected screen, the current screen content, followed
 963         by differences between the two screens.
 964         """
 965
 966         assert self._mode == self.MODE_TEST
 967
 968         # Print the expected screen. The output is not verbatim as it was
 969         # specified in the input file, but instead the screen is printed in the
 970         # same way the current screen gets output. This allows to properly show
 971         # differences between the two screens.
 972
 973         expected_screen_e = self._get_screen_xml(expected_screen)
 974         self._indent_xml(expected_screen_e)
 975         expected_screen_str = ElementTree.tostring(
 976             expected_screen_e, 'unicode')
 977         print("Expected (normalized) screen:", file=sys.stderr)
 978         print(expected_screen_str, file=sys.stderr)
 979
 980         # Print the current screen.
 981         current_screen_e = self._get_screen_xml(self._screen)
 982         self._indent_xml(current_screen_e)
 983         current_screen_str = ElementTree.tostring(current_screen_e, 'unicode')
 984         print("Current screen:", file=sys.stderr)
 985         print(current_screen_str, file=sys.stderr)
 986
 987         # Print the delta.
 988         print("Differences:", file=sys.stderr)
 989         sys.stderr.writelines(difflib.unified_diff(
 990             expected_screen_str.splitlines(keepends=True),
 991             current_screen_str.splitlines(keepends=True),
 992             fromfile="Expected screen", tofile="Current screen"))
 993
 994     def _execute_playbook(self, test_e):
 995         """
 996         Run the main loop and execute the given test playbook. Normal return
 997         from the method indicates that the test succeeded. Exception
 998         _TestFailure is raised when the test fails and exception
 999         _TerminalConnectionException can be thrown when communication with the
1000         pseudo-terminal fails.
1001         """
1002
1003         assert self._mode == self.MODE_TEST
1004
1005         cmd_iter = iter(test_e)
1006
1007         # Start the main loop.
1008         with selectors.DefaultSelector() as sel:
1009             sel.register(self._fd, selectors.EVENT_READ)
1010
1011             expected_screen = None
1012             more_commands = True
1013             while True:
1014                 # Process any actions and find an expected screen.
1015                 while not expected_screen and more_commands:
1016                     try:
1017                         cmd_e = next(cmd_iter)
1018                         if cmd_e.tag == 'action':
1019                             self._playbook_key(cmd_e)
1020                         elif cmd_e.tag == 'expect':
1021                             expected_screen = self._parse_expected_screen(
1022                                 cmd_e)
1023                             # Stop processing more commands for now and wait
1024                             # for the expected screen to appear.
1025                             break
1026                         else:
1027                             raise self._TestFailure(
1028                                 "Element '{}' is invalid, expected 'action' "
1029                                 "or 'expect'".format(cmd_e.tag))
1030                     except StopIteration:
1031                         # No more commands.
1032                         more_commands = False
1033
1034                 # Wait for the expected screen.
1035                 events = sel.select(CHILD_TIMEOUT)
1036                 if not events:
1037                     if expected_screen:
1038                         self._report_failed_expectation(expected_screen)
1039                     raise self._TestFailure(
1040                         "Timeout reached. No event received in the last {} "
1041                         "second(s)".format(CHILD_TIMEOUT))
1042
1043                 # Expect only an event on self._fd.
1044                 assert len(events) == 1
1045                 event = events[0]
1046                 key, _mask = event
1047                 assert key.fd == self._fd
1048
1049                 closed = self._pty_callback()
1050                 if closed:
1051                     if more_commands:
1052                         raise self._TestFailure(
1053                             "Connection to the terminal was closed but the "
1054                             "playbook contains more commands")
1055                     break
1056
1057                 # Check if the expected screen is present.
1058                 if self._screen == expected_screen:
1059                     expected_screen = None
1060
1061     def execute_test(self, filename):
1062         """
1063         Load test data from a given file, start the program under the test and
1064         execute the test playbook. Returns True when the test succeeded, False
1065         otherwise.
1066         """
1067
1068         assert self._mode == self.MODE_TEST
1069
1070         # Read the test data.
1071         try:
1072             tree = ElementTree.ElementTree(file=filename)
1073         except Exception as e:
1074             print("Failed to read playbook file '{}': {}.".format(filename, e),
1075                   file=sys.stderr)
1076             return False
1077
1078         # Read what program to execute.
1079         test_e = tree.getroot()
1080         if test_e.tag != 'test':
1081             print("Root element '{}' is invalid, expected 'test'.".format(
1082                   test_e.tag), file=sys.stderr)
1083             return False
1084         try:
1085             self._program = test_e.attrib['program']
1086         except KeyError:
1087             print("Element 'test' is missing required attribute 'program'.",
1088                   file=sys.stderr)
1089             return False
1090
1091         # Start the specified program.
1092         if not self._start_program():
1093             return False
1094
1095         # Execute the test playbook.
1096         res = True
1097         try:
1098             self._execute_playbook(tree.getroot())
1099         except (self._TerminalConnectionException, self._TestFailure) as e:
1100             print("{}.".format(e), file=sys.stderr)
1101             res = False
1102         finally:
1103             # Finalize the run of the child program.
1104             if not self._finalize_program():
1105                 res = False
1106
1107         # Return whether the test passed.
1108         return res
1109
1110
1111 def main():
1112     """
1113     Parse command line arguments and execute the operation that the user
1114     selected. Returns 0 if the operation was successful and a non-zero value
1115     otherwise.
1116     """
1117
1118     # Parse command line arguments.
1119     parser = argparse.ArgumentParser()
1120     parser.add_argument(
1121         '-t', '--terminfo', metavar='PATH', help="path to terminfo directory")
1122
1123     subparsers = parser.add_subparsers(dest='program')
1124     subparsers.required = True
1125
1126     program_parser = argparse.ArgumentParser(add_help=False)
1127     program_parser.add_argument('program', help="executable to run")
1128
1129     # Create the parser for the 'run' command.
1130     parser_run = subparsers.add_parser(
1131         'run', parents=[program_parser], help="run a program")
1132     parser_run.set_defaults(mode=Term.MODE_RUN)
1133
1134     # Create the parser for the 'record' command.
1135     parser_record = subparsers.add_parser(
1136         'record', parents=[program_parser], help="record a test")
1137     parser_record.set_defaults(mode=Term.MODE_RECORD)
1138     parser_record.add_argument(
1139         '-o', '--playbook', metavar='FILE', required=True,
1140         help="output playbook file")
1141
1142     # Create the parser for the 'test' command.
1143     parser_test = subparsers.add_parser('test', help="perform a test")
1144     parser_test.set_defaults(program=None)
1145     parser_test.set_defaults(mode=Term.MODE_TEST)
1146     parser_test.add_argument('playbook', help="input playbook file")
1147
1148     args = parser.parse_args()
1149
1150     tk_root = None
1151     if args.mode in (Term.MODE_RUN, Term.MODE_RECORD):
1152         # Start the terminal GUI.
1153         global tkinter
1154         import tkinter
1155
1156         try:
1157             tk_root = tkinter.Tk()
1158         except tkinter.TclError as e:
1159             print("Failed to initialize GUI: {}.".format(e), file=sys.stderr)
1160             return 1
1161
1162     term = Term(tk_root, args.program, args.mode, args.terminfo)
1163     if tk_root:
1164         # Start the GUI main loop.
1165         term.run_gui_mainloop()
1166     else:
1167         # Execute and check the playbook, without running GUI.
1168         ok = term.execute_test(args.playbook)
1169         if ok:
1170             msg = "succeeded"
1171             res = 0
1172         else:
1173             msg = "failed"
1174             res = 1
1175
1176         print("Checking of playbook '{}' {}.".format(args.playbook, msg))
1177         return res
1178
1179     if args.mode == Term.MODE_RECORD:
1180         # Get the recorded test data and write them to a file.
1181         if not term.output_test(args.playbook):
1182             return 1
1183
1184     return 0
1185
1186
1187 if __name__ == '__main__':
1188     sys.exit(main())
1189
1190 # vim: set tabstop=4 shiftwidth=4 textwidth=79 expandtab :