Python: Refactor field-flow test

python/ql/test/experimental/dataflow/fieldflow/test.py

@@ -27,7 +27,10 @@ def SINK_F(x):
         print("OK")
 
 
-# Preamble
+# ------------------------------------------------------------------------------
+# Actual tests
+# ------------------------------------------------------------------------------
+
 class MyObj(object):
     def __init__(self, foo):
         self.foo = foo
@@ -35,61 +38,31 @@ class MyObj(object):
     def setFoo(self, foo):
         self.foo = foo
 
-
-class NestedObj(object):
-    def __init__(self):
-        self.obj = MyObj("OK")
-
-    def getObj(self):
-        return self.obj
-
-# Actual tests
-
 def setFoo(obj, x):
     SINK_F(obj.foo)
     obj.foo = x
 
 @expects(2)
-def test_example1():
+def test_indirect_assign():
     myobj = MyObj("OK")
 
     setFoo(myobj, SOURCE)
     SINK(myobj.foo) # $ flow="SOURCE, l:-1 -> myobj.foo"
 
 
-def test_example1_method():
+def test_indirect_assign_method():
     myobj = MyObj("OK")
 
     myobj.setFoo(SOURCE)
     SINK(myobj.foo) # $ MISSING: flow
 
 
-def test_example2():
-    x = SOURCE
-
-    a = NestedObj()
-
-    a.obj.foo = x
-
-    SINK(a.obj.foo) # $ flow="SOURCE, l:-6 -> a.obj.foo"
-
-
-def test_example2_method():
-    x = SOURCE
-
-    a = NestedObj()
-
-    a.getObj().foo = x
-
-    SINK(a.obj.foo) # $ MISSING: flow
-
-
-def test_example3():
+def test_constructor_assign():
     obj = MyObj(SOURCE)
     SINK(obj.foo) # $ flow="SOURCE, l:-1 -> obj.foo"
 
 
-def test_example3_kw():
+def test_constructor_assign_kw():
     obj = MyObj(foo=SOURCE)
     SINK(obj.foo) # $ flow="SOURCE, l:-1 -> obj.foo"
 
@@ -99,6 +72,30 @@ def fields_with_local_flow(x):
     a = obj.foo
     return a
 
-
 def test_fields():
     SINK(fields_with_local_flow(SOURCE)) # $ flow="SOURCE -> fields_with_local_flow(..)"
+
+# ------------------------------------------------------------------------------
+# Nested Object
+# ------------------------------------------------------------------------------
+
+class NestedObj(object):
+    def __init__(self):
+        self.obj = MyObj("OK")
+
+    def getObj(self):
+        return self.obj
+
+
+def test_nested_obj():
+    x = SOURCE
+    a = NestedObj()
+    a.obj.foo = x
+    SINK(a.obj.foo) # $ flow="SOURCE, l:-3 -> a.obj.foo"
+
+
+def test_nested_obj_method():
+    x = SOURCE
+    a = NestedObj()
+    a.getObj().foo = x
+    SINK(a.obj.foo) # $ MISSING: flow