Python: split test as suggested in review

python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/test_for.py

@@ -0,0 +1,30 @@
+# Add taintlib to PATH so it can be imported during runtime without any hassle
+import sys; import os; sys.path.append(os.path.dirname(os.path.dirname((__file__))))
+from taintlib import *
+
+# This has no runtime impact, but allows autocomplete to work
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from ..taintlib import *
+
+
+# Actual tests
+
+class Iter:
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        raise StopIteration
+
+def test_for():
+    iter = Iter()
+    taint(iter)
+    for tainted in iter:
+        ensure_tainted(tainted) # $ tainted
+
+
+
+# Make tests runable
+
+test_for()

python/ql/test/experimental/dataflow/tainttracking/defaultAdditionalTaintStep/test_with.py

@@ -51,19 +51,6 @@ def test_with_arg():
     with ctx as tainted:
         ensure_tainted(tainted) # $ tainted
 
-class Iter:
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        raise StopIteration
-
-def test_for():
-    iter = Iter()
-    taint(iter)
-    for tainted in iter:
-        ensure_tainted(tainted) # $ tainted
-
 
 
 # Make tests runable
@@ -71,4 +58,3 @@ def test_for():
 test_with()
 test_with_taint()
 test_with_arg()
-test_for()