Python: Improve tests of json module

Inspired by the work on previous commit

Author: RasmusWL

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

@@ -11,54 +11,43 @@
 # Actual tests
 
 from io import StringIO
-
-# Workaround for Python3 not having unicode
-import sys
-if sys.version_info[0] == 3:
-    unicode = str
+import json
 
 def test():
     print("\n# test")
     ts = TAINTED_STRING
-    import json
+
+    encoded = json.dumps(ts)
 
     ensure_tainted(
+        encoded, # $ tainted
         json.dumps(ts), # $ tainted
-        json.loads(json.dumps(ts)), # $ tainted
+        json.dumps(obj=ts), # $ MISSING: tainted
+        json.loads(encoded), # $ tainted
+        json.loads(s=encoded), # $ MISSING: tainted
     )
 
-    # For Python2, need to convert to unicode for StringIO to work
-    tainted_filelike = StringIO(unicode(json.dumps(ts)))
+    # load/dump with file-like
+    tainted_filelike = StringIO()
+    json.dump(ts, tainted_filelike)
 
+    tainted_filelike.seek(0)
     ensure_tainted(
         tainted_filelike, # $ MISSING: tainted
         json.load(tainted_filelike), # $ MISSING: tainted
     )
 
-def non_syntacical():
-    print("\n# non_syntacical")
-    ts = TAINTED_STRING
-
-    # a less syntactical approach
-    from json import load, loads, dumps
-
-    dumps_alias = dumps
-
-    ensure_tainted(
-        dumps(ts), # $ tainted
-        dumps_alias(ts), # $ tainted
-        loads(dumps(ts)), # $ tainted
-    )
-
-    # For Python2, need to convert to unicode for StringIO to work
-    tainted_filelike = StringIO(unicode(dumps(ts)))
+    # load/dump with file-like using keyword-args
+    tainted_filelike = StringIO()
+    json.dump(obj=ts, fp=tainted_filelike)
 
+    tainted_filelike.seek(0)
     ensure_tainted(
         tainted_filelike, # $ MISSING: tainted
-        load(tainted_filelike), # $ MISSING: tainted
+        json.load(fp=tainted_filelike), # $ MISSING: tainted
     )
 
+
 # Make tests runable
 
 test()
-non_syntacical()