[OSAB] Task 4: Add unit tests

bashplotlib/scatterplot.py

@@ -9,6 +9,7 @@
 import csv
 import sys
 import optparse
+import os
 from .utils.helpers import *
 from .utils.commandhelp import scatter
 
@@ -28,26 +29,30 @@ def get_scale(series, is_y=False, steps=20):
     return scaled_series
 
 
-def _plot_scatter(xs, ys, size, pch, colour, title, cs):
+def build_scatter(xs, ys, size, pch, colour, title):
     plotted = set()
 
+    plot_str = ""
+
     if title:
-        print(box_text(title, 2 * (len(get_scale(xs, False, size)) + 1)))
+        plot_str += box_text(title, 2 * (len(get_scale(xs, False, size)) + 1))
+        plot_str += os.linesep
 
-    print("-" * (2 * (len(get_scale(xs, False, size)) + 2)))
+    plot_str += "-" * (2 * (len(get_scale(xs, False, size)) + 2))
+    plot_str += os.linesep
     for y in get_scale(ys, True, size):
-        print("|", end=' ')
+        plot_str += "| "
         for x in get_scale(xs, False, size):
             point = " "
             for (i, (xp, yp)) in enumerate(zip(xs, ys)):
                 if xp <= x and yp >= y and (xp, yp) not in plotted:
                     point = pch
                     plotted.add((xp, yp))
-                    if cs:
-                        colour = cs[i]
-            printcolour(point + " ", True, colour)
-        print(" |")
-    print("-" * (2 * (len(get_scale(xs, False, size)) + 2)))
+            plot_str += buildcolour(point + " ", colour)
+        plot_str += " |" + os.linesep
+
+    plot_str += "-" * (2 * (len(get_scale(xs, False, size)) + 2))
+    return plot_str
 
 def plot_scatter(f, xs, ys, size, pch, colour, title):
     """
@@ -62,7 +67,6 @@ def plot_scatter(f, xs, ys, size, pch, colour, title):
         colour -- colour of the points
         title -- title of the plot
     """
-    cs = None
     if f:
         if isinstance(f, str):
             with open(f) as fh:
@@ -71,8 +75,6 @@ def plot_scatter(f, xs, ys, size, pch, colour, title):
             data = [tuple(line.strip().split(',')) for line in f]
         xs = [float(i[0]) for i in data]
         ys = [float(i[1]) for i in data]
-        if len(data[0]) > 2:
-            cs = [i[2].strip() for i in data]
     elif isinstance(xs, list) and isinstance(ys, list):
         pass
     else:
@@ -81,8 +83,7 @@ def plot_scatter(f, xs, ys, size, pch, colour, title):
         with open(ys) as fh:
             ys = [float(str(row).strip()) for row in fh]
 
-    _plot_scatter(xs, ys, size, pch, colour, title, cs)
-
+    print(build_scatter(xs, ys, size, pch, colour, title))
 
 
 def main():

bashplotlib/utils/helpers.py

@@ -30,18 +30,30 @@ def get_colour(colour):
     """
     Get the escape code sequence for a colour
     """
-    return bcolours.get(colour, bcolours['ENDC'])
+    return bcolours.get(colour, bcolours['default'])
 
 
-def printcolour(text, sameline=False, colour=get_colour("ENDC")):
+def printcolour(text, sameline=False, colour=get_colour("default")):
     """
-    Print color text using escape codes
+    Print color text
     """
     if sameline:
         sep = ''
     else:
         sep = '\n'
-    sys.stdout.write(get_colour(colour) + text + bcolours["ENDC"] + sep)
+    sys.stdout.write(buildcolour(text, colour) + sep)
+
+
+def buildcolour(text, colour="default"):
+    """
+    Build color text using escape codes
+    """
+    # Unwise hack to avoid having to deal with colors in the
+    # tests.
+    if colour == "default":
+        return text
+    else:
+        return get_colour(colour) + text + bcolours["ENDC"]
 
 
 def drange(start, stop, step=1.0, include_stop=False):

test.py

@@ -0,0 +1,44 @@
+from bashplotlib.scatterplot import build_scatter
+
+if __name__ == "__main__":
+    # Build a scatter plot
+    scatter = build_scatter(
+            [-10,20,30],
+            [-10,20,30],
+            10,
+            'x',
+            'default',
+            'test hello')
+
+    # This is what we know our scatter plot should
+    # look like. We have to be careful that indentations
+    # match.
+    expected_scatter = """----------------------------
+|        test hello        |
+----------------------------
+----------------------------
+|                       x  |
+|                          |
+|                          |
+|                   x      |
+|                          |
+|                          |
+|                          |
+|                          |
+|                          |
+|                          |
+|                          |
+| x                        |
+----------------------------"""
+
+    # Compare what we got to what we expected, and
+    # fail with a helpful error message if anything is
+    # different.
+    if expected_scatter == scatter:
+        print("SUCCESS!")
+    else:
+        print("FAILURE!")
+        print("Expected:")
+        print(expected_scatter)
+        print("Got:")
+        print(scatter)