gistfile1.diff

diff --git a/lib/riot/context.rb b/lib/riot/context.rb
index 8532ccd..324f395 100644
--- a/lib/riot/context.rb
+++ b/lib/riot/context.rb
@@ -1,5 +1,10 @@
 module Riot
-  RootContext = Struct.new(:setups, :teardowns)
+  RootContext = Struct.new(:setups, :teardowns) do
+    def assertion_class
+      Assertion
+    end
+  end 
+
   class Context
     attr_reader :description
     def initialize(description, parent=RootContext.new([],[]), &definition)
@@ -23,13 +28,25 @@ module Riot
       @contexts << self.class.new("#{@description} #{description}", self, &definition)
     end
     
+    def extend_assertions(*extension_modules)
+      @assertion_class = Class.new(Assertion) do
+        include *extension_modules
+      end
+    end
+
     def run(reporter)
       reporter.describe_context(self) unless @assertions.empty?
       local_run(reporter, Situation.new)
       run_sub_contexts(reporter)
       reporter
     end
+
+    def assertion_class
+      @assertion_class ||= @parent.assertion_class
+    end
+
   private
+
     def local_run(reporter, situation)
       (setups + @assertions + teardowns).each do |runnable|
         reporter.report(runnable.to_s, runnable.run(situation))
@@ -39,7 +56,7 @@ module Riot
     def run_sub_contexts(reporter) @contexts.each { |ctx| ctx.run(reporter) }; end
 
     def new_assertion(scope, what, &definition)
-      (@assertions << Assertion.new("#{scope} #{what}", &definition)).last
+      (@assertions << assertion_class.new("#{scope} #{what}", &definition)).last
     end
   end # Context
 end # Riot
diff --git a/test/context_test.rb b/test/context_test.rb
index 7605a69..50a4618 100644
--- a/test/context_test.rb
+++ b/test/context_test.rb
@@ -88,3 +88,55 @@ context "The asserts_topic shortcut" do
     topic.equals("bar").run(situation)
   end.equals([:pass])
 end # The asserts_topic shortcut
+
+context "Setting assertion extensions" do
+  fake_mod_one, fake_mod_two = Module.new, Module.new
+
+  setup do
+    Riot::Context.new("foo") do
+      extend_assertions fake_mod_one, fake_mod_two
+    end.asserts_topic
+  end
+
+  should("still return an Assertion") { topic }.kind_of(Riot::Assertion)
+
+  should("have included the given assertion extension modules in the assertion") do
+    topic.class.included_modules.include?(fake_mod_one) && topic.class.included_modules.include?(fake_mod_two)
+  end
+
+  context "involving subcontexts without the subcontext extending assertions" do
+    assertion_one = assertion_two = nil
+
+    setup do 
+      Riot::Context.new "bleh" do
+        extend_assertions Module.new
+        assertion_one = asserts_topic
+        context("foo") { assertion_two = asserts_topic } 
+      end
+    end
+    
+    should "not create separate instances of the assertion class in subcontexts" do
+      assertion_one && assertion_two && assertion_one.class.object_id == assertion_two.class.object_id
+    end
+  end # involving subcontexts without the subcontext extending assertions
+
+  context "involving subcontexts with the subcontext extending assertions" do
+    assertion_one = assertion_two = nil
+    
+    setup do
+      Riot::Context.new "bah" do
+        extend_assertions Module.new
+        assertion_one = asserts_topic
+        context("meh") do
+          extend_assertions Module.new
+          assertion_two = asserts_topic
+        end
+      end
+    end
+
+    should "create separate instances of the assertion class in subcontexts" do
+      assertion_one && assertion_two && assertion_one.class.object_id != assertion_two.class.object_id
+    end
+
+  end # involving subcontexts with the subcontext extending assertions
+end   # Setting assertion extensions