From 7ff35cc7b20f712057a8fd84abbe28229542b468 Mon Sep 17 00:00:00 2001
From: Dan Kubb <dan.kubb@gmail.com>
Date: Sun, 28 Jul 2013 15:14:39 -0700
Subject: [PATCH] Fix the mutator registry to be more strict about nodes

---
 lib/mutant/mutator/registry.rb | 58 +++++++++++++++++++++++++++-------
 1 file changed, 47 insertions(+), 11 deletions(-)

diff --git a/lib/mutant/mutator/registry.rb b/lib/mutant/mutator/registry.rb
index 157a576b..afd5c0d0 100644
--- a/lib/mutant/mutator/registry.rb
+++ b/lib/mutant/mutator/registry.rb
@@ -2,16 +2,12 @@ module Mutant
   class Mutator
     # Registry for mutators
     module Registry
-      # Return registry state
-      #
-      # @return [Hash]
-      #
-      # @api private
-      #
-      def self.registry
-        @registry ||= {}
-      end
-      private_class_method :registry
+
+      # Raised when the type is an invalid type
+      InvalidTypeError = Class.new(TypeError)
+
+      # Raised when the type is a duplicate
+      DuplicateTypeError = Class.new(ArgumentError)
 
       # Register mutator class for AST node class
       #
@@ -23,7 +19,8 @@ module Mutant
       # @return [self]
       #
       def self.register(type, mutator_class)
-        raise "duplicate type registration: #{type}" if registry.key?(type)
+        assert_valid_type(type)
+        assert_unique_type(type)
         registry[type] = mutator_class
         self
       end
@@ -46,6 +43,45 @@ module Mutant
         end
       end
 
+      # Return registry state
+      #
+      # @return [Hash]
+      #
+      # @api private
+      #
+      def self.registry
+        @registry ||= {}
+      end
+      private_class_method :registry
+
+      # Assert the node type is valid
+      #
+      # @raise [InvalidTypeError]
+      #   raised when the node type is invalid
+      #
+      # @api private
+      #
+      def self.assert_valid_type(type)
+        unless NODE_TYPES.include?(type) || type.kind_of?(Class)
+          raise InvalidTypeError, "invalid type registration: #{type}"
+        end
+      end
+      private_class_method :assert_valid_type
+
+      # Assert the node type is unique and not already registered
+      #
+      # @raise [DuplcateTypeError]
+      #   raised when the node type is a duplicate
+      #
+      # @api private
+      #
+      def self.assert_unique_type(type)
+        if registry.key?(type)
+          raise DuplicateTypeError, "duplicate type registration: #{type}"
+        end
+      end
+      private_class_method :assert_unique_type
+
     end # Registry
   end # Mutator
 end # Mutant