Add a method to fetch the schema for a document

This will prevent having to pass a schema object around all the time

lib/engines/content_block_manager/app/models/content_block_manager/content_block/document.rb

@@ -50,6 +50,10 @@ def latest_draft
         editions.where(state: :draft).order(created_at: :asc).last
       end
 
+      def schema
+        @schema ||= ContentBlockManager::ContentBlock::Schema.find_by_block_type(block_type)
+      end
+
     private
 
       def embed_code_prefix

lib/engines/content_block_manager/test/factories/content_block_document.rb

@@ -7,10 +7,19 @@
     latest_edition_id { nil }
     live_edition_id { nil }
 
+    transient do
+      schema { nil }
+    end
+
     ContentBlockManager::ContentBlock::Schema.valid_schemas.each do |type|
       trait type.to_sym do
         block_type { type }
+        schema { build(:content_block_schema, block_type: type) }
       end
     end
+
+    after(:build) do |content_block_document, evaluator|
+      content_block_document.stubs(:schema).returns(evaluator.schema)
+    end
   end
 end

lib/engines/content_block_manager/test/unit/app/models/content_block_document_test.rb

@@ -185,4 +185,20 @@ class ContentBlockManager::ContentBlockDocumentTest < ActiveSupport::TestCase
       assert_equal newest_draft, document.latest_draft
     end
   end
+
+  describe "#schema" do
+    let(:document) { build(:content_block_document, :email_address) }
+    let(:schema) { build(:content_block_schema) }
+
+    it "returns a schema object" do
+      document.unstub(:schema)
+
+      ContentBlockManager::ContentBlock::Schema
+        .expects(:find_by_block_type)
+        .with(document.block_type)
+        .returns(schema)
+
+      assert_equal document.schema, schema
+    end
+  end
 end