diff --git a/.gitignore b/.gitignore
index 48f84b9..fb6abe9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -17,3 +17,5 @@ spec/support/active_record.rb
 test/tmp
 test/version_tmp
 tmp
+.byebug_history
+polymorphic_integer_type_test
diff --git a/lib/polymorphic_integer_type.rb b/lib/polymorphic_integer_type.rb
index ef0a12c..e2e0fec 100644
--- a/lib/polymorphic_integer_type.rb
+++ b/lib/polymorphic_integer_type.rb
@@ -1,5 +1,6 @@
 require "polymorphic_integer_type/version"
 require "polymorphic_integer_type/extensions"
 require "polymorphic_integer_type/mapping"
+require "polymorphic_integer_type/polymorphic_array_value_extension"
 
 module PolymorphicIntegerType; end
diff --git a/lib/polymorphic_integer_type/polymorphic_array_value_extension.rb b/lib/polymorphic_integer_type/polymorphic_array_value_extension.rb
new file mode 100644
index 0000000..0e14a98
--- /dev/null
+++ b/lib/polymorphic_integer_type/polymorphic_array_value_extension.rb
@@ -0,0 +1,20 @@
+module PolymorphicIntegerType
+  module PolymorphicArrayValueExtension
+    def type_to_ids_mapping
+      super.tap do |result|
+        association = @associated_table.send(:association)
+        klass = association.active_record
+        name = association.name
+
+        if klass.respond_to?("#{name}_type_mapping")
+          result.transform_keys! do |key|
+            klass.send("#{name}_type_mapping").key(key)
+          end
+        end
+        result
+      end
+    end
+  end
+end
+
+ActiveRecord::PredicateBuilder::PolymorphicArrayValue.prepend(PolymorphicIntegerType::PolymorphicArrayValueExtension)
diff --git a/lib/polymorphic_integer_type/version.rb b/lib/polymorphic_integer_type/version.rb
index bac1c0d..0c8bb9c 100644
--- a/lib/polymorphic_integer_type/version.rb
+++ b/lib/polymorphic_integer_type/version.rb
@@ -1,3 +1,3 @@
 module PolymorphicIntegerType
-  VERSION = "2.1.0"
+  VERSION = "2.2.0"
 end
diff --git a/polymorphic_integer_type.gemspec b/polymorphic_integer_type.gemspec
index 67a4ab1..e91540d 100644
--- a/polymorphic_integer_type.gemspec
+++ b/polymorphic_integer_type.gemspec
@@ -18,10 +18,10 @@ Gem::Specification.new do |spec|
   spec.test_files    = spec.files.grep(%r{^(test|spec|features)/})
   spec.require_paths = ["lib"]
 
-  spec.add_development_dependency "bundler", "~> 1.2"
+  spec.add_dependency "activerecord"
+  spec.add_development_dependency "bundler"
   spec.add_development_dependency "rake"
   spec.add_development_dependency "rspec"
-  spec.add_development_dependency "activerecord", "4.2.0"
-  spec.add_development_dependency "mysql2", "0.3.16"
+  spec.add_development_dependency "sqlite3"
   spec.add_development_dependency "byebug"
 end
diff --git a/spec/polymorphic_integer_type_spec.rb b/spec/polymorphic_integer_type_spec.rb
index a5ab798..0516f32 100644
--- a/spec/polymorphic_integer_type_spec.rb
+++ b/spec/polymorphic_integer_type_spec.rb
@@ -24,6 +24,17 @@
     end
   end
 
+  context "when querying the associations" do
+    let(:source) { cat }
+    let(:target) { nil }
+    it "properly finds the object with a where" do
+      expect(Link.where(source: source, id: link.id).first).to eql link
+    end
+    it "properly finds the object with a find_by" do
+      expect(Link.find_by(source: source, id: link.id)).to eql link
+    end
+  end
+
   shared_examples "proper source" do
     it "should have the proper id, type and object for the source" do
       expect(link.source_id).to eql source.id
@@ -47,12 +58,12 @@
 
     context "and the link is accessed through the associations" do
       before { link }
-      
+
       it "should have the proper source" do
         expect(source.source_links[0].source).to eql source
       end
     end
-    
+
   end
   context "When a link is given polymorphic record" do
     let(:link) { Link.create(source: source) }
@@ -137,7 +148,7 @@
 
 
   end
-  
+
   context "when the association is an STI table" do
     let(:link) { Link.create(source: source, target: whiskey) }
     let(:source) { Dog.create(name: "Bela", kind: "Dog", owner: owner) }
@@ -147,7 +158,7 @@
       expect(link.source).to eql source
     end
   end
-  
+
   context "when mapping is given inline in the belongs_to model" do
     class InlineLink < ActiveRecord::Base
       include PolymorphicIntegerType::Extensions
@@ -187,7 +198,7 @@ class InlineDrink < ActiveRecord::Base
     include_examples "proper target"
 
     it "creates foreign_type mapping method" do
-      expect(Link.source_type_mapping).to eq({0 => "Person", 1 => "Animal"})
+      expect(Link.source_type_mapping).to eq({1 => "Person", 2 => "Animal"})
       expect(InlineLink.source_type_mapping).to eq({10 => "Person", 11 => "InlineAnimal"})
     end
 
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 7bb3688..9523223 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -13,6 +13,7 @@
   config.before(:suite) do
     database_config = YAML.load(File.open("#{File.dirname(__FILE__)}/support/database.yml"))
     ActiveRecord::Base.establish_connection(database_config)
+    ActiveRecord::MigrationContext.new("#{File.dirname(__FILE__)}/support/migrations").migrate
   end
 
   config.around do |example|
diff --git a/spec/support/configuration.rb b/spec/support/configuration.rb
index be58de9..2a6cfd8 100644
--- a/spec/support/configuration.rb
+++ b/spec/support/configuration.rb
@@ -1,4 +1,4 @@
 PolymorphicIntegerType::Mapping.configuration do |config|
-  config.add :source, {0 => "Person", 1 => "Animal"}
-  config.add :target, {0 => "Food", 1 => "Drink"}
+  config.add :source, {1 => "Person", 2 => "Animal"}
+  config.add :target, {1 => "Food", 2 => "Drink"}
 end
diff --git a/spec/support/database.yml b/spec/support/database.yml
index 581336c..ca63670 100644
--- a/spec/support/database.yml
+++ b/spec/support/database.yml
@@ -1,5 +1,5 @@
 ---
-:adapter: mysql2
+:adapter: sqlite3
 :host: localhost
 :database: polymorphic_integer_type_test
 :username: root
diff --git a/spec/support/migrations/1_create_link_table.rb b/spec/support/migrations/1_create_link_table.rb
index b1d0c84..2518e95 100644
--- a/spec/support/migrations/1_create_link_table.rb
+++ b/spec/support/migrations/1_create_link_table.rb
@@ -1,4 +1,4 @@
-class CreateLinkTable < ActiveRecord::Migration
+class CreateLinkTable < ActiveRecord::Migration[5.2]
 
   def up
     create_table :links do |t|
diff --git a/spec/support/migrations/2_create_animal_table.rb b/spec/support/migrations/2_create_animal_table.rb
index e2038d4..7341dbd 100644
--- a/spec/support/migrations/2_create_animal_table.rb
+++ b/spec/support/migrations/2_create_animal_table.rb
@@ -1,4 +1,4 @@
-class CreateAnimalTable < ActiveRecord::Migration
+class CreateAnimalTable < ActiveRecord::Migration[5.2]
 
   def up
     create_table :animals do |t|
diff --git a/spec/support/migrations/3_create_person_table.rb b/spec/support/migrations/3_create_person_table.rb
index 64a2414..5103dde 100644
--- a/spec/support/migrations/3_create_person_table.rb
+++ b/spec/support/migrations/3_create_person_table.rb
@@ -1,4 +1,4 @@
-class CreatePersonTable < ActiveRecord::Migration
+class CreatePersonTable < ActiveRecord::Migration[5.2]
 
   def up
     create_table :people do |t|
diff --git a/spec/support/migrations/4_create_food_table.rb b/spec/support/migrations/4_create_food_table.rb
index 728deb3..5dbb37a 100644
--- a/spec/support/migrations/4_create_food_table.rb
+++ b/spec/support/migrations/4_create_food_table.rb
@@ -1,4 +1,4 @@
-class CreateFoodTable < ActiveRecord::Migration
+class CreateFoodTable < ActiveRecord::Migration[5.2]
 
   def up
     create_table :foods do |t|
diff --git a/spec/support/migrations/5_create_drink_table.rb b/spec/support/migrations/5_create_drink_table.rb
index 36dc26e..bdcd807 100644
--- a/spec/support/migrations/5_create_drink_table.rb
+++ b/spec/support/migrations/5_create_drink_table.rb
@@ -1,4 +1,4 @@
-class CreateDrinkTable < ActiveRecord::Migration
+class CreateDrinkTable < ActiveRecord::Migration[5.2]
 
   def up
     create_table :drinks do |t|