lib/flipper/registry.rb in flipper-0.3.0 vs lib/flipper/registry.rb in flipper-0.4.0

@@ -1,15 +1,25 @@
 require 'thread'
 
 module Flipper
+  # Internal: Used to store registry of groups by name.
   class Registry
     include Enumerable
 
     class Error < StandardError; end
     class DuplicateKey < Error; end
-    class MissingKey < Error; end
 
+    class KeyNotFound < Error
+      # Public: The key that was not found
+      attr_reader :key
+
+      def initialize(key)
+        @key = key
+        super("Key #{key.inspect} not found")
+      end
+    end
+
     def initialize(source = {})
       @mutex = Mutex.new
       @source = source
     end
 
@@ -20,22 +30,28 @@
     def values
       @mutex.synchronize { @source.values }
     end
 
     def add(key, value)
-      @mutex.synchronize do
+      key = key.to_sym
+
+      @mutex.synchronize {
         if @source[key]
           raise DuplicateKey, "#{key} is already registered"
         else
           @source[key] = value
         end
-      end
+      }
     end
 
     def get(key)
-      @mutex.synchronize do
-        @source[key]
-      end
+      key = key.to_sym
+
+      @mutex.synchronize {
+        @source.fetch(key) {
+          raise KeyNotFound.new(key)
+        }
+      }
     end
 
     def each(&block)
       @mutex.synchronize { @source.dup }.each(&block)
     end