faif · faif · Jan 13, 2019 · Jan 12, 2019 · Jan 12, 2019
diff --git a/append_output.sh b/append_output.sh
@@ -1,12 +1,19 @@
 #!/bin/bash
 
+# This script (given path to a python script as an argument)
+# appends python outputs to given file.
+
 set -e
 
-src=$(sed -n -e '/### OUTPUT ###/,$!p' "$1")
-output=$(python "$1" | sed 's/^/# /')
+output_marker='OUTPUT = """'
+
+# get everything (excluding part between `output_marker` and the end of the file)
+# into `src` var
+src=$(sed -n -e "/$output_marker/,\$!p" "$1")
+output=$(python "$1")
 
-# These are done separately to avoid having to insert a newline, which causes
-# problems when the text itself has '\n' in strings
 echo "$src" > $1
-echo -e "\n### OUTPUT ###" >> $1
+echo -e "\n" >> $1
+echo "$output_marker" >> $1
 echo "$output" >> $1
+echo '"""' >> $1
diff --git a/behavioral/strategy.py b/behavioral/strategy.py
@@ -36,15 +36,21 @@ def on_sale_discount(order):
     return order.price * 0.25 + 20
 
 
-if __name__ == "__main__":
+def main():
     order0 = Order(100)
     order1 = Order(100, discount_strategy=ten_percent_discount)
     order2 = Order(1000, discount_strategy=on_sale_discount)
     print(order0)
     print(order1)
     print(order2)
 
-### OUTPUT ###
-# <Price: 100, price after discount: 100>
-# <Price: 100, price after discount: 90.0>
-# <Price: 1000, price after discount: 730.0>
+
+if __name__ == "__main__":
+    main()
+
+
+OUTPUT = """
+<Price: 100, price after discount: 100>
+<Price: 100, price after discount: 90.0>
+<Price: 1000, price after discount: 730.0>
+"""
diff --git a/behavioral/visitor.py b/behavioral/visitor.py
@@ -45,15 +45,22 @@ def visit_B(self, node, *args, **kwargs):
         print('visit_B ' + node.__class__.__name__)
 
 
-a = A()
-b = B()
-c = C()
-visitor = Visitor()
-visitor.visit(a)
-visitor.visit(b)
-visitor.visit(c)
-
-### OUTPUT ###
-# generic_visit A
-# visit_B B
-# visit_B C
+def main():
+    a = A()
+    b = B()
+    c = C()
+    visitor = Visitor()
+    visitor.visit(a)
+    visitor.visit(b)
+    visitor.visit(c)
+
+
+if __name__ == "__main__":
+    main()
+
+
+OUTPUT = """
+generic_visit A
+visit_B B
+visit_B C
+"""
diff --git a/tests/test_outputs.py b/tests/test_outputs.py
@@ -0,0 +1,29 @@
+try:
+    from contextlib import redirect_stdout
+except:
+    pass
+
+import io
+import sys
+
+import pytest
+
+from behavioral.visitor import main as visitor_main
+from behavioral.visitor import OUTPUT as visitor_output
+from behavioral.strategy import main as strategy_main
+from behavioral.strategy import OUTPUT as strategy_output
+
+@pytest.mark.skipif(sys.version_info < (3,4),
+                    reason="requires python3.4 or higher")
+@pytest.mark.parametrize("main,output", [
+    (visitor_main, visitor_output),
+    (strategy_main, strategy_output),
+])
+def test_output(main, output):
+    f = io.StringIO()
+    with redirect_stdout(f):
+        main()
+
+    real_output = f.getvalue().strip()
+    expected_output = output.strip()
+    assert real_output == expected_output
diff --git a/tests/test_strategy.py b/tests/test_strategy.py