Add more ptop tests

px/px_top.py

@@ -126,6 +126,14 @@ def get_toplist(baseline):
     return toplist
 
 
+def writebytes(bytestring):
+    if sys.version_info.major == 2:
+        sys.stdout.write(bytestring)
+    else:
+        # http://stackoverflow.com/a/908440/473672
+        sys.stdout.buffer.write(bytestring)
+
+
 def clear_screen():
     """
     Clear the screen and move cursor to top left corner:
@@ -139,41 +147,43 @@ def clear_screen():
     """
 
     CSI = b"\x1b["
-    if sys.version_info.major == 2:
-        sys.stdout.write(CSI + b"1J")
-        sys.stdout.write(CSI + b"H")
-    else:
-        # http://stackoverflow.com/a/908440/473672
-        sys.stdout.buffer.write(CSI + b"1J")
-        sys.stdout.buffer.write(CSI + b"H")
-    sys.stdout.flush()
+    writebytes(CSI + b"1J")
+    writebytes(CSI + b"H")
 
 
-def redraw(load_bar, baseline, rows, columns):
-    """
-    Refresh display relative to the given baseline.
-
-    The new display will be (at most) rows rows x columns columns.
-    """
+def get_screen_lines(load_bar, baseline, rows, columns):
     load = px_load.get_load_values()
-    loadstring = px_load.get_load_string(load)
+    loadstring = px_load.get_load_string(load).encode('utf-8')
     loadbar = load_bar.get_bar(load=load[0], columns=40)
     lines = [
-        "System load: " + loadstring + " |{}|".format(loadbar),
-        ""]
+        b"System load: " + loadstring + b" |" + loadbar + b"|",
+        b""]
 
     toplist_table_lines = px_terminal.to_screen_lines(get_toplist(baseline), columns)
     if toplist_table_lines:
         heading_line = toplist_table_lines[0]
         heading_line = px_terminal.get_string_of_length(heading_line, columns)
         heading_line = px_terminal.inverse_video(heading_line)
         toplist_table_lines[0] = heading_line
+
+    toplist_table_lines = map(lambda s: s.encode('utf-8'), toplist_table_lines)
     lines += toplist_table_lines
 
-    clear_screen()
+    return lines[0:rows]
+
+
+def redraw(load_bar, baseline, rows, columns, clear=True):
+    """
+    Refresh display relative to the given baseline.
+
+    The new display will be (at most) rows rows x columns columns.
+    """
+    lines = get_screen_lines(load_bar, baseline, rows, columns)
+    if clear:
+        clear_screen()
 
     # We need both \r and \n when the TTY is in tty.setraw() mode
-    sys.stdout.write("\r\n".join(lines[0:rows]))
+    writebytes(b"\r\n".join(lines))
     sys.stdout.flush()
 
 

tests/px_top_test.py

@@ -74,4 +74,29 @@ def test_redraw():
     # Just make sure it doesn't crash
     loadbar = px_load_bar.PxLoadBar(1, 1)
     baseline = px_process.get_all()
-    px_top.redraw(loadbar, baseline, 100, 10)
+    px_top.redraw(loadbar, baseline, 100, 10, clear=False)
+
+
+def test_get_screen_lines():
+    loadbar = px_load_bar.PxLoadBar(1, 1)
+    baseline = px_process.get_all()
+
+    SCREEN_ROWS = 10
+    SCREEN_COLUMNS = 100
+    lines = px_top.get_screen_lines(
+        loadbar, baseline, SCREEN_ROWS, SCREEN_COLUMNS)
+
+    # Top row should contain ANSI escape codes
+    CSI = b"\x1b["
+    assert b'CSI' in lines[0].replace(CSI, b'CSI')
+
+    assert len(lines) == SCREEN_ROWS
+
+    # This is a bit fluffy, but due to ANSI escape codes the line lengths are
+    # difficult to measure. So if at least one line is exactly the right length
+    # then we're probably fine.
+    found_right_length = False
+    for line in lines:
+        if len(line) == SCREEN_COLUMNS:
+            found_right_length = True
+    assert found_right_length