Fix object handling in yaml for PyYAML 4.x

docs/source/loader.rst

@@ -13,13 +13,14 @@ build and run tests with pytest_ with some limitations described below.
 .. note:: It is also possible to run gabbi tests from the command
           line. See :doc:`runner`.
 
-.. note:: By default gabbi will load YAML files using the ``safe_load``
-          function. This means only basic YAML types are allowed in the
-          file. For most use cases this is fine. If you need custom types
-          (for example, to match NaN) it is possible to set the ``safe_yaml``
-          parameter of :meth:`~gabbi.driver.build_tests` to ``False``.
-          If custom types are used, please keep in mind that this can limit
-          the portability of the YAML files to other contexts.
+.. note:: By default gabbi will load YAML files in the default "safe"
+          fashion. This means that arbitrary Python objects cannot be
+          loaded. Custom YAML types which are defined in the process may be
+          used. If for some reason you need unsafe support, you may
+          set the ``safe_yaml`` parameter of
+          :meth:`~gabbi.driver.build_tests` to ``False``.
+          Keep in mind that if custom types are used, this can limit the
+          portability of the YAML files to other contexts.
 
 .. warning:: If test are being run with a runner that supports
              concurrency (such as ``testrepository``) it is critical

gabbi/tests/gabbits_intercept/nan-safe.yaml

@@ -0,0 +1,9 @@
+tests:
+  - name: test NAN tag
+    url: /nan
+    method: GET
+    request_headers:
+        content-type: application/json
+    response_json_paths:
+        $.nan: !NanChecker {}
+        $.nan: !IsNAN

gabbi/tests/gabbits_runner/nan-safe.yaml

@@ -0,0 +1,8 @@
+tests:
+  - name: test NAN
+    url: /nan
+    method: GET
+    request_headers:
+        content-type: application/json
+    response_json_paths:
+        $.nan: !NanChecker {}

gabbi/tests/gabbits_unsafe_yaml/nan.yaml

@@ -1,9 +1,8 @@
 tests:
-  - name: test NAN
+  - name: test NAN object
     url: /nan
     method: GET
     request_headers:
         content-type: application/json
     response_json_paths:
-        $.nan: !NanChecker {}
-        $.nan: !IsNAN
+        $.nan: !!python/object:gabbi.tests.util.NanChecker {}

gabbi/tests/test_runner.py

@@ -16,6 +16,7 @@
 import sys
 import unittest
 from uuid import uuid4
+import yaml
 
 from six import StringIO
 from wsgi_intercept.interceptor import Urllib3Interceptor
@@ -25,6 +26,8 @@
 from gabbi.handlers.jsonhandler import JSONHandler
 from gabbi import runner
 from gabbi.tests.simple_wsgi import SimpleWsgi
+# To register yaml additions.
+import gabbi.tests.util  # noqa
 
 
 class RunnerTest(unittest.TestCase):
@@ -85,19 +88,48 @@ def test_input_files(self):
             except SystemExit as err:
                 self.assertFailure(err)
 
-    def test_unsafe_yaml(self):
+    def test_unsafe_yaml_success(self):
         sys.argv = ['gabbi-run', 'http://%s:%s/nan' % (self.host, self.port)]
 
         sys.argv.append('--unsafe-yaml')
         sys.argv.append('--')
-        sys.argv.append('gabbi/tests/gabbits_runner/nan.yaml')
+        sys.argv.append('gabbi/tests/gabbits_runner/nan-unsafe.yaml')
 
         with self.server():
             try:
                 runner.run()
             except SystemExit as err:
                 self.assertSuccess(err)
 
+    def test_unsafe_yaml_failure(self):
+        sys.argv = ['gabbi-run', 'http://%s:%s/nan' % (self.host, self.port)]
+
+        sys.argv.append('--')
+        sys.argv.append('gabbi/tests/gabbits_runner/nan-unsafe.yaml')
+
+        self.assertRaises(yaml.constructor.ConstructorError, runner.run)
+
+    def test_safe_yaml_success(self):
+        sys.argv = ['gabbi-run', 'http://%s:%s/nan' % (self.host, self.port)]
+
+        sys.argv.append('--')
+        sys.argv.append('gabbi/tests/gabbits_runner/nan-safe.yaml')
+
+        with self.server():
+            try:
+                runner.run()
+            except SystemExit as err:
+                self.assertSuccess(err)
+
+    def test_safe_yaml_failure(self):
+        sys.argv = ['gabbi-run', 'http://%s:%s/nan' % (self.host, self.port)]
+
+        sys.argv.append('--unsafe-yaml')
+        sys.argv.append('--')
+        sys.argv.append('gabbi/tests/gabbits_runner/nan-safe.yaml')
+
+        self.assertRaises(yaml.constructor.ConstructorError, runner.run)
+
     def test_target_url_parsing(self):
         sys.argv = ['gabbi-run', 'http://%s:%s/foo' % (self.host, self.port)]
 

gabbi/tests/test_unsafe_yaml.py

@@ -18,23 +18,18 @@
 
 import os
 
-import yaml
-
 from gabbi import driver
 # TODO(cdent): test_pytest allows pytest to see the tests this module
 # produces. Without it, the generator will not run. It is a todo because
 # needing to do this is annoying and gross.
 from gabbi.driver import test_pytest  # noqa
 from gabbi.tests import simple_wsgi
-from gabbi.tests import util
+from gabbi.tests import util  # noqa
 
 
 TESTS_DIR = 'gabbits_unsafe_yaml'
 
 
-yaml.add_constructor(u'!IsNAN', lambda loader, node: util.NanChecker())
-
-
 BUILD_TEST_ARGS = dict(
     intercept=simple_wsgi.SimpleWsgi,
     safe_yaml=False

gabbi/tests/util.py

@@ -40,3 +40,5 @@ def __eq__(self, other):
             return math.isnan(other)
         except TypeError:
             return False
+
+yaml.add_constructor(u'!IsNAN', lambda loader, node: NanChecker())

gabbi/utils.py

@@ -104,7 +104,13 @@ def load_yaml(handle=None, yaml_file=None, safe=True):
 
     If no file or handle is provided, read from STDIN.
     """
-    load = yaml.safe_load if safe else yaml.load
+    # In PyYAML 4.1 safe is the default
+    try:
+        danger_load = yaml.danger_load
+    except AttributeError:
+        danger_load = yaml.load
+    load = yaml.safe_load if safe else danger_load
+
 
     if yaml_file:
         with io.open(yaml_file, encoding='utf-8') as source: