diff --git a/tests/demo1_correct/code/tree.py b/tests/demo1_correct/code/tree.py
index f4ad1f7467fcdc48c1708adc9cc99911d57e88c0..cf12a7bdf0e8e0ba367deacaa1770f9400af8b9e 100644
--- a/tests/demo1_correct/code/tree.py
+++ b/tests/demo1_correct/code/tree.py
@@ -5,7 +5,6 @@ class TreeNode(object):
         # TODO: 3 lines missing.
         raise NotImplementedError("Insert your solution and remove this error.")
 
-#!s=a
 def sorted_array_to_bst(nums):
     # TODO: 4 lines missing.
     raise NotImplementedError("Your stuff here")
@@ -15,7 +14,6 @@ def sorted_array_to_bst(nums):
     raise NotImplementedError("Here")
 print("hello world asdfasd")
 
-#!o=a
 def preOrder(node):
     if not node: 
         # TODO: 1 lines missing.
@@ -25,21 +23,16 @@ def preOrder(node):
     preOrder(node.right)
 for _ in range(10):
     print("Hello world")
-#!o=a
 a = 234
 
 result = sorted_array_to_bst([1, 2, 3, 4, 5, 6, 7])
 preOrder(result)
 
-#!i=a
 for _ in range(10):
     print("hi")
-#!i=a
 
-#!i=b
 print("hello")
 def myfun(a):
     return a*2
 
 print(myfun(4))
-#!i=b
diff --git a/tests/demo1_correct/output/tree_myfile.py b/tests/demo1_correct/output/tree_myfile.py
index 900f89bbfda7f14940aa056d74482ee49e77c072..d06d5ffbf51dcd0d4d8543e682c00a3c6a6276f5 100644
--- a/tests/demo1_correct/output/tree_myfile.py
+++ b/tests/demo1_correct/output/tree_myfile.py
@@ -7,7 +7,6 @@
     return node 
 print("hello world asdfasd")
 
-#!o=a
 def preOrder(node):
     if not node: 
         return
diff --git a/tests/demo1_correct/output/tree_myfile_stripped.py b/tests/demo1_correct/output/tree_myfile_stripped.py
index 900f89bbfda7f14940aa056d74482ee49e77c072..d06d5ffbf51dcd0d4d8543e682c00a3c6a6276f5 100644
--- a/tests/demo1_correct/output/tree_myfile_stripped.py
+++ b/tests/demo1_correct/output/tree_myfile_stripped.py
@@ -7,7 +7,6 @@
     return node 
 print("hello world asdfasd")
 
-#!o=a
 def preOrder(node):
     if not node: 
         return
diff --git a/tests/setup_test_files.py b/tests/setup_test_files.py
index 4181e127ed9ee6e6342f27f833f423518a66b3fc..918f4ef6a307b0425e0ada426957df98a6e680b6 100644
--- a/tests/setup_test_files.py
+++ b/tests/setup_test_files.py
@@ -1,5 +1,6 @@
 import os.path
 import shutil
+dir = os.path.dirname(__file__)
 
 
 def setup(source, dest):
@@ -33,8 +34,8 @@ if __name__ == "__main__":
     # if not os.path.isdir("student_correct_output"):
     #     os.mkdir("student_correct_output")
     for file in ['demo1']:
-        setup(file, file+"_correct")
-    setup_keep("demo2/framework.py", "demo2/framework_correct.txt")
+        setup(dir+"/" + file, dir+"/" + file+"_correct")
+    setup_keep(dir+"/demo2/framework.py", dir+"/demo2/framework_correct.txt")
 
     a= 234
     # assert False