rb: Support `raise Xception, message` for Structs that inherit from ::Exception [THRIFT-58]

Author: Kevin Ballard <kevin@rapleaf.com>


git-svn-id: https://svn.apache.org/repos/asf/incubator/thrift/trunk@680542 13f79535-47bb-0310-9956-ffa450edef68
diff --git a/lib/rb/lib/thrift/struct.rb b/lib/rb/lib/thrift/struct.rb
index 10f2ea0..d3de200 100644
--- a/lib/rb/lib/thrift/struct.rb
+++ b/lib/rb/lib/thrift/struct.rb
@@ -80,6 +80,32 @@
 
     protected
 
+    def self.append_features(mod)
+      if mod.ancestors.include? ::Exception
+        mod.send :class_variable_set, :'@@__thrift_struct_real_initialize', mod.instance_method(:initialize)
+        super
+        # set up our custom initializer so `raise Xception, 'message'` works
+        mod.send :define_method, :struct_initialize, mod.instance_method(:initialize)
+        mod.send :define_method, :initialize, mod.instance_method(:exception_initialize)
+      else
+        super
+      end
+    end
+
+    def exception_initialize(*args, &block)
+      if args.size == 1 and args.first.is_a? Hash
+        # looks like it's a regular Struct initialize
+        method(:struct_initialize).call(args.first)
+      else
+        # call the Struct initializer first with no args
+        # this will set our field default values
+        method(:struct_initialize).call()
+        # now give it to the exception
+        self.class.send(:class_variable_get, :'@@__thrift_struct_real_initialize').bind(self).call(*args, &block)
+        # self.class.instance_method(:initialize).bind(self).call(*args, &block)
+      end
+    end
+
     def handle_message(iprot, fid, ftype)
       field = struct_fields[fid]
       if field and field[:type] == ftype
diff --git a/lib/rb/spec/struct_spec.rb b/lib/rb/spec/struct_spec.rb
index 2245778..15e8437 100644
--- a/lib/rb/spec/struct_spec.rb
+++ b/lib/rb/spec/struct_spec.rb
@@ -5,11 +5,12 @@
   include Thrift
   include SpecNamespace
 
-  class OldStruct
+  class Xception < Thrift::Exception
     include Thrift::Struct
-    attr_accessor :set
+    attr_accessor :message, :code
     FIELDS = {
-      1 => {:type => Thrift::Types::SET, :name => 'val', :default => {:foo => true, :bar => true}}
+      1 => {:type => Thrift::Types::STRING, :name => 'message'},
+      2 => {:type => Thrift::Types::I32, :name => 'code', :default => 1}
     }
   end
 
@@ -199,5 +200,40 @@
       lambda { Hello.new(:fish => 'salmon') }.should raise_error(Exception, "Unknown keys given to SpecNamespace::Hello.new: fish")
       lambda { Hello.new(:foo => 'bar', :baz => 'blah', :greeting => 'Good day') }.should raise_error(Exception, /^Unknown keys given to SpecNamespace::Hello.new: (foo, baz|baz, foo)$/)
     end
+
+    it "should support `raise Xception, 'message'` for Exception structs" do
+      begin
+        raise Xception, "something happened"
+      rescue Thrift::Exception => e
+        e.message.should == "something happened"
+        e.code.should == 1
+        # ensure it gets serialized properly, this is the really important part
+        prot = mock("Protocol")
+        prot.should_receive(:write_struct_begin).with("ThriftStructSpec::Xception")
+        prot.should_receive(:write_struct_end)
+        prot.should_receive(:write_field).with('message', Types::STRING, 1, "something happened")
+        prot.should_receive(:write_field).with('code', Types::I32, 2, 1)
+        prot.should_receive(:write_field_stop)
+
+        e.write(prot)
+      end
+    end
+
+    it "should support the regular initializer for exception structs" do
+      begin
+        raise Xception, :message => "something happened", :code => 5
+      rescue Thrift::Exception => e
+        e.message.should == "something happened"
+        e.code.should == 5
+        prot = mock("Protocol")
+        prot.should_receive(:write_struct_begin).with("ThriftStructSpec::Xception")
+        prot.should_receive(:write_struct_end)
+        prot.should_receive(:write_field).with('message', Types::STRING, 1, "something happened")
+        prot.should_receive(:write_field).with('code', Types::I32, 2, 5)
+        prot.should_receive(:write_field_stop)
+
+        e.write(prot)
+      end
+    end
   end
 end