1 files changed, 22 insertions, 20 deletions
diff --git a/_test/lib/test_tokens.py b/_test/lib/test_tokens.py
index 1483db8..575e95c 100644
--- a/_test/lib/test_tokens.py
+++ b/_test/lib/test_tokens.py
@@ -1,5 +1,5 @@
 
-import ruamel.yaml as yaml
+import ruamel.yaml
 import pprint
 
 # Tokens mnemonic:
@@ -22,24 +22,24 @@ import pprint
 # value:                :
 
 _replaces = {
-    yaml.DirectiveToken: '%',
-    yaml.DocumentStartToken: '---',
-    yaml.DocumentEndToken: '...',
-    yaml.AliasToken: '*',
-    yaml.AnchorToken: '&',
-    yaml.TagToken: '!',
-    yaml.ScalarToken: '_',
-    yaml.BlockSequenceStartToken: '[[',
-    yaml.BlockMappingStartToken: '{{',
-    yaml.BlockEndToken: ']}',
-    yaml.FlowSequenceStartToken: '[',
-    yaml.FlowSequenceEndToken: ']',
-    yaml.FlowMappingStartToken: '{',
-    yaml.FlowMappingEndToken: '}',
-    yaml.BlockEntryToken: ',',
-    yaml.FlowEntryToken: ',',
-    yaml.KeyToken: '?',
-    yaml.ValueToken: ':',
+    ruamel.yaml.DirectiveToken: '%',
+    ruamel.yaml.DocumentStartToken: '---',
+    ruamel.yaml.DocumentEndToken: '...',
+    ruamel.yaml.AliasToken: '*',
+    ruamel.yaml.AnchorToken: '&',
+    ruamel.yaml.TagToken: '!',
+    ruamel.yaml.ScalarToken: '_',
+    ruamel.yaml.BlockSequenceStartToken: '[[',
+    ruamel.yaml.BlockMappingStartToken: '{{',
+    ruamel.yaml.BlockEndToken: ']}',
+    ruamel.yaml.FlowSequenceStartToken: '[',
+    ruamel.yaml.FlowSequenceEndToken: ']',
+    ruamel.yaml.FlowMappingStartToken: '{',
+    ruamel.yaml.FlowMappingEndToken: '}',
+    ruamel.yaml.BlockEntryToken: ',',
+    ruamel.yaml.FlowEntryToken: ',',
+    ruamel.yaml.KeyToken: '?',
+    ruamel.yaml.ValueToken: ':',
 }
 
 
@@ -48,9 +48,10 @@ def test_tokens(data_filename, tokens_filename, verbose=False):
     with open(tokens_filename, 'r') as fp:
         tokens2 = fp.read().split()
     try:
+        yaml = ruamel.yaml.YAML(typ='unsafe', pure=True)
         with open(data_filename, 'rb') as fp1:
             for token in yaml.scan(fp1):
-                if not isinstance(token, (yaml.StreamStartToken, yaml.StreamEndToken)):
+                if not isinstance(token, (ruamel.yaml.StreamStartToken, ruamel.yaml.StreamEndToken)):
                     tokens1.append(_replaces[token.__class__])
     finally:
         if verbose:
@@ -68,6 +69,7 @@ def test_scanner(data_filename, canonical_filename, verbose=False):
     for filename in [data_filename, canonical_filename]:
         tokens = []
         try:
+            yaml = ruamel.yaml.YAML(typ='unsafe', pure=False)
             with open(filename, 'rb') as fp:
                 for token in yaml.scan(fp):
                     tokens.append(token.__class__.__name__)