1 files changed, 23 insertions, 17 deletions
diff --git a/test/psych/helper.rb b/test/psych/helper.rb
index 9348457958..4e82887c6d 100644
--- a/test/psych/helper.rb
+++ b/test/psych/helper.rb
@@ -1,5 +1,5 @@
 # frozen_string_literal: true
-require 'minitest/autorun'
+require 'test/unit'
 require 'stringio'
 require 'tempfile'
 require 'date'
@@ -7,13 +7,7 @@ require 'date'
 require 'psych'
 
 module Psych
-  superclass = if defined?(Minitest::Test)
-                 Minitest::Test
-               else
-                 MiniTest::Unit::TestCase
-               end
-
-  class TestCase < superclass
+  class TestCase < Test::Unit::TestCase
     def self.suppress_warning
       verbose, $VERBOSE = $VERBOSE, nil
       yield
@@ -47,24 +41,30 @@ module Psych
     # Convert between Psych and the object to verify correct parsing and
     # emitting
     #
-    def assert_to_yaml( obj, yaml )
-      assert_equal( obj, Psych::load( yaml ) )
+    def assert_to_yaml( obj, yaml, loader = :load )
+      assert_equal( obj, Psych.send(loader, yaml) )
       assert_equal( obj, Psych::parse( yaml ).transform )
-      assert_equal( obj, Psych::load( obj.to_yaml ) )
+      assert_equal( obj, Psych.send(loader, obj.to_yaml) )
       assert_equal( obj, Psych::parse( obj.to_yaml ).transform )
-      assert_equal( obj, Psych::load(
+      assert_equal( obj, Psych.send(loader,
         obj.to_yaml(
           :UseVersion => true, :UseHeader => true, :SortKeys => true
         )
       ))
+    rescue Psych::DisallowedClass, Psych::BadAlias, Psych::AliasesNotEnabled
+      assert_to_yaml obj, yaml, :unsafe_load
     end
 
     #
     # Test parser only
     #
     def assert_parse_only( obj, yaml )
-      assert_equal( obj, Psych::load( yaml ) )
-      assert_equal( obj, Psych::parse( yaml ).transform )
+      begin
+        assert_equal obj, Psych::load( yaml )
+      rescue Psych::DisallowedClass, Psych::BadAlias, Psych::AliasesNotEnabled
+        assert_equal obj, Psych::unsafe_load( yaml )
+      end
+      assert_equal obj, Psych::parse( yaml ).transform
     end
 
     def assert_cycle( obj )
@@ -75,9 +75,15 @@ module Psych
         assert_nil Psych::load(Psych.dump(obj))
         assert_nil Psych::load(obj.to_yaml)
       else
-        assert_equal(obj, Psych.load(v.tree.yaml))
-        assert_equal(obj, Psych::load(Psych.dump(obj)))
-        assert_equal(obj, Psych::load(obj.to_yaml))
+        begin
+          assert_equal(obj, Psych.load(v.tree.yaml))
+          assert_equal(obj, Psych::load(Psych.dump(obj)))
+          assert_equal(obj, Psych::load(obj.to_yaml))
+        rescue Psych::DisallowedClass, Psych::BadAlias, Psych::AliasesNotEnabled
+          assert_equal(obj, Psych.unsafe_load(v.tree.yaml))
+          assert_equal(obj, Psych::unsafe_load(Psych.dump(obj)))
+          assert_equal(obj, Psych::unsafe_load(obj.to_yaml))
+        end
       end
     end