address review.

code_to_optimize/bubble_sort_method.py

@@ -1,3 +1,6 @@
+import sys
+
+
 class BubbleSorter:
     def __init__(self, x=0):
         self.x = x
@@ -10,4 +13,5 @@ class BubbleSorter:
                     temp = arr[j]
                     arr[j] = arr[j + 1]
                     arr[j + 1] = temp
+        print("stderr test", file=sys.stderr)
         return arr

tests/test_instrument_all_and_run.py

@@ -168,8 +168,17 @@ def test_sort():
             pytest_max_loops=1,
             testing_time=0.1,
         )
-        assert "codeflash stdout: Sorting list" in test_results[0].stdout
-        assert "result: [0, 1, 2, 3, 4, 5]" in test_results[0].stdout
+
+        out_str = """--------------------------------- Captured Log ---------------------------------
+
+--------------------------------- Captured Out ---------------------------------
+
+codeflash stdout: Sorting list
+result: [0, 1, 2, 3, 4, 5]
+
+codeflash stdout: Sorting list
+result: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]"""
+        assert out_str == test_results[0].stdout.strip()
         assert test_results[0].id.function_getting_tested == "sorter"
         assert test_results[0].id.iteration_id == "1_0"
         assert test_results[0].id.test_class_name is None
@@ -181,8 +190,8 @@ def test_sort():
         assert test_results[0].runtime > 0
         assert test_results[0].did_pass
         assert test_results[0].return_value == ([0, 1, 2, 3, 4, 5],)
-        assert "codeflash stdout: Sorting list" in test_results[1].stdout
-        assert "result: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]" in test_results[1].stdout
+        assert out_str == test_results[1].stdout.strip()
+
         assert test_results[1].id.function_getting_tested == "sorter"
         assert test_results[1].id.iteration_id == "4_0"
         assert test_results[1].id.test_class_name is None

tests/test_instrument_tests.py

@@ -482,8 +482,16 @@ def test_sort():
         )
         assert test_results_perf[1].runtime > 0
         assert test_results_perf[1].did_pass
-        assert "codeflash stdout: Sorting list" in test_results_perf[1].stdout
-        assert "result: [0, 1, 2, 3, 4, 5]" in test_results_perf[1].stdout
+        out_str = """--------------------------------- Captured Log ---------------------------------
+
+--------------------------------- Captured Out ---------------------------------
+codeflash stdout: Sorting list
+result: [0, 1, 2, 3, 4, 5]
+
+codeflash stdout: Sorting list
+result: [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]"""
+        assert out_str == test_results_perf[1].stdout
+
     finally:
         test_path.unlink(missing_ok=True)
         test_path_perf.unlink(missing_ok=True)

tests/test_instrumentation_run_results_aiservice.py

@@ -187,6 +187,9 @@ def test_single_element_list():
 
         # Replace with optimized code that mutated instance attribute
         optimized_code_mutated_attr = """
+import sys
+
+
 class BubbleSorter:
 
     def __init__(self, x=1):
@@ -200,6 +203,7 @@ class BubbleSorter:
                     temp = arr[j]
                     arr[j] = arr[j + 1]
                     arr[j + 1] = temp
+        print("stderr test", file=sys.stderr)
         return arr
                         """
         fto_path.write_text(optimized_code_mutated_attr, "utf-8")
@@ -332,6 +336,9 @@ def test_single_element_list():
         assert test_results[1].return_value[2] == [1, 2, 3]
         # Replace with optimized code that mutated instance attribute
         optimized_code_mutated_attr = """
+import sys
+
+
 class BubbleSorter:
 
     def __init__(self, x=1):
@@ -345,6 +352,7 @@ class BubbleSorter:
                     temp = arr[j]
                     arr[j] = arr[j + 1]
                     arr[j + 1] = temp
+        print("stderr test", file=sys.stderr)
         return arr
                         """
         fto_path.write_text(optimized_code_mutated_attr, "utf-8")
@@ -388,6 +396,9 @@ class BubbleSorter:
         )  # The test should fail because the instance attribute was mutated
         # Replace with optimized code that did not mutate existing instance attribute, but added a new one
         optimized_code_new_attr = """
+import sys
+
+
 class BubbleSorter:
     def __init__(self, x=0):
         self.x = x
@@ -401,6 +412,7 @@ class BubbleSorter:
                     temp = arr[j]
                     arr[j] = arr[j + 1]
                     arr[j + 1] = temp
+        print("stderr test", file=sys.stderr)
         return arr
                         """
         fto_path.write_text(optimized_code_new_attr, "utf-8")