Fix Custom Strategy

lib/polo/configuration.rb

@@ -6,11 +6,28 @@ class Configuration
     def initialize(options={})
       options = { on_duplicate: nil, obfuscate: {} }.merge(options)
       @on_duplicate_strategy = options[:on_duplicate]
-      @blacklist = options[:obfuscate]
+      obfuscate(options[:obfuscate])
     end
 
+    # TODO: document this
+    # This normalizes an array or hash of fields to a hash of
+    # { field_name => strategy }
     def obfuscate(*fields)
-      @blacklist = fields
+      if fields.is_a?(Array)
+        fields = fields.flatten
+      end
+
+      fields_and_strategies = {}
+
+      fields.each do |field|
+        if field.is_a?(Symbol) || field.is_a?(String)
+          fields_and_strategies[field] = nil
+        elsif field.is_a?(Hash)
+          fields_and_strategies = fields_and_strategies.merge(field)
+        end
+      end
+
+      @blacklist = fields_and_strategies
     end
 
     def on_duplicate(strategy)

lib/polo/translator.rb

@@ -36,6 +36,7 @@ def instances
     def obfuscate!(instances, fields)
       instances.each do |instance|
         next if intersection(instance.attributes.keys, fields).empty?
+
         fields.each do |field, strategy|
           value = instance.attributes[field.to_s] || ''
           instance.send("#{field}=", new_field_value(field, strategy, value))
@@ -44,7 +45,7 @@ def obfuscate!(instances, fields)
     end
 
     def intersection(attrs, fields)
-      attrs & fields.to_a.flatten.map(&:to_s)
+      attrs & fields.keys.map(&:to_s)
     end
 
     def new_field_value(field, strategy, value)

spec/configuration_spec.rb

@@ -28,18 +28,20 @@
       end
 
       defaults = Polo.defaults
-      expect(defaults.blacklist).to eq([:email, :password])
+      expect(defaults.blacklist).to eq({ :email => nil, :password => nil })
     end
 
     it 'allows for the user to define fields with strategies in blacklist' do
       Polo.configure do
+
         email_strategy = lambda {|e| "#{e.split("@")[0]}_test@example.com" }
         credit_card_strategy = lambda {|_| "4111111111111111"}
+
         obfuscate({email: email_strategy, credit_card: credit_card_strategy})
       end
 
       defaults = Polo.defaults
-      expect(defaults.blacklist.first[:email].call("a@b.com")).to eq("a_test@example.com")
+      expect(defaults.blacklist[:email].call("a@b.com")).to eq("a_test@example.com")
     end
   end
 end

spec/polo_spec.rb

@@ -78,6 +78,16 @@
         expect(scrambled_email).to_not eq('nettofarah@gmail.com')
         expect(insert).to match(exp.first)
       end
+
+      it 'can apply custom strategies' do
+        Polo.configure do
+          obfuscate(email: lambda { |_| 'changeme' })
+        end
+
+        inserts = Polo.explore(AR::Chef, 1)
+
+        expect(inserts).to eq [ %q{INSERT INTO `chefs` (`id`, `name`, `email`) VALUES (1, 'Netto', 'changeme')} ]
+      end
     end
 
     describe 'on_duplicate' do