diff --git a/lib/diaspora_federation/test/factories.rb b/lib/diaspora_federation/test/factories.rb
index 253d73f..3f3809b 100644
--- a/lib/diaspora_federation/test/factories.rb
+++ b/lib/diaspora_federation/test/factories.rb
@@ -41,7 +41,7 @@ module DiasporaFederation
 
       Fabricator(:account_migration_entity, class_name: DiasporaFederation::Entities::AccountMigration) do
         author { Fabricate.sequence(:diaspora_id) }
-        profile { Fabricate(:profile_entity) }
+        profile {|attrs| Fabricate(:profile_entity, author: attrs[:author]) }
         old_identity { Fabricate.sequence(:diaspora_id) }
       end
 
diff --git a/lib/diaspora_federation/validators/optional_aware_validator.rb b/lib/diaspora_federation/validators/optional_aware_validator.rb
index f0df658..1c0b467 100644
--- a/lib/diaspora_federation/validators/optional_aware_validator.rb
+++ b/lib/diaspora_federation/validators/optional_aware_validator.rb
@@ -13,10 +13,8 @@ module DiasporaFederation
       private
 
       def optional_props
-        entity_name = self.class.name.split("::").last.sub("Validator", "")
-        return [] unless Entities.const_defined?(entity_name)
-
-        Entities.const_get(entity_name).optional_props
+        return [] unless @obj.class.respond_to?(:optional_props)
+        @obj.class.optional_props
       end
     end
   end
diff --git a/spec/lib/diaspora_federation/validators/optional_aware_validator_spec.rb b/spec/lib/diaspora_federation/validators/optional_aware_validator_spec.rb
index 164675a..049ee49 100644
--- a/spec/lib/diaspora_federation/validators/optional_aware_validator_spec.rb
+++ b/spec/lib/diaspora_federation/validators/optional_aware_validator_spec.rb
@@ -1,47 +1,46 @@
 module DiasporaFederation
   describe Validators::OptionalAwareValidator do
-    let(:entity_data) {
-      {test1: "abc", test2: true, test7: "abc", multi: []}
-    }
+    def entity_stub(additional_data={})
+      allow_any_instance_of(Entities::TestComplexEntity).to receive(:freeze)
+      allow_any_instance_of(Entities::TestComplexEntity).to receive(:validate)
+      entity_data = {test1: "abc", test2: true, test3: nil, test4: nil, test5: nil, test6: nil, test7: "abc", multi: []}
+      Entities::TestComplexEntity.new(entity_data.merge(additional_data))
+    end
 
     it "validates a valid object" do
-      validator = Validators::TestComplexEntityValidator.new(OpenStruct.new(entity_data))
+      validator = Validators::TestComplexEntityValidator.new(entity_stub)
       expect(validator).to be_valid
       expect(validator.errors).to be_empty
     end
 
     it "fails when a mandatory property is invalid" do
       ["ab", nil].each do |val|
-        entity = OpenStruct.new(entity_data.merge(test1: val))
-        validator = Validators::TestComplexEntityValidator.new(entity)
+        validator = Validators::TestComplexEntityValidator.new(entity_stub(test1: val))
         expect(validator).not_to be_valid
         expect(validator.errors).to include(:test1)
       end
     end
 
     it "fails when an optional property is invalid" do
-      entity = OpenStruct.new(entity_data.merge(test7: "ab"))
-      validator = Validators::TestComplexEntityValidator.new(entity)
+      validator = Validators::TestComplexEntityValidator.new(entity_stub(test7: "ab"))
       expect(validator).not_to be_valid
       expect(validator.errors).to include(:test7)
     end
 
     it "allows an optional property to be nil" do
-      entity = OpenStruct.new(entity_data.merge(test7: nil))
-      validator = Validators::TestComplexEntityValidator.new(entity)
+      validator = Validators::TestComplexEntityValidator.new(entity_stub(test7: nil))
       expect(validator).to be_valid
       expect(validator.errors).to be_empty
     end
 
     it "doesn't ignore 'not_nil' rules for an optional property" do
-      entity = OpenStruct.new(entity_data.merge(multi: nil))
-      validator = Validators::TestComplexEntityValidator.new(entity)
+      validator = Validators::TestComplexEntityValidator.new(entity_stub(multi: nil))
       expect(validator).not_to be_valid
       expect(validator.errors).to include(:multi)
     end
 
-    it "doesn't fail when there is no entity for this validator" do
-      entity = OpenStruct.new(entity_data.merge(test1: nil))
+    it "doesn't fail when the entity doesn't have optional props" do
+      entity = OpenStruct.new(test1: nil)
       validator = Validators::TestUnknownEntityValidator.new(entity)
       expect(validator).not_to be_valid
       expect(validator.errors).to include(:test1)
diff --git a/spec/support/shared_validator_specs.rb b/spec/support/shared_validator_specs.rb
index 16def0c..97f46ee 100644
--- a/spec/support/shared_validator_specs.rb
+++ b/spec/support/shared_validator_specs.rb
@@ -1,6 +1,8 @@
 def entity_stub(entity, data={})
-  OpenStruct.new(Fabricate.schematic(entity).options[:class_name].default_values
-                   .merge(Fabricate.attributes_for(entity)).merge(data))
+  entity_class = Fabricate.schematic(entity).options[:class_name]
+  allow_any_instance_of(entity_class).to receive(:freeze)
+  allow_any_instance_of(entity_class).to receive(:validate)
+  entity_class.new(Fabricate.attributes_for(entity).merge(data))
 end
 
 ALPHANUMERIC_RANGE = [*"0".."9", *"A".."Z", *"a".."z"].freeze