From afb82d2a5a1f3f08cca7fa6ec7e7201dd455695f Mon Sep 17 00:00:00 2001
From: Nick Thomas <nick@lupine.me.uk>
Date: Fri, 10 Jun 2011 23:01:00 +0100
Subject: [PATCH] Add tests for the IPv6Address class and ensure they pass

---
 lib/btc_wire_proto.rb       | 14 +++++++---
 spec/btc_wire_proto_spec.rb | 56 +++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/lib/btc_wire_proto.rb b/lib/btc_wire_proto.rb
index e54cd5f..8474063 100644
--- a/lib/btc_wire_proto.rb
+++ b/lib/btc_wire_proto.rb
@@ -22,19 +22,25 @@ class Ipv6Address < BinData::Primitive
   end
 
   def set(v)
+    v = IPAddress(v) if v.is_a?(String)
     v = IPAddress::IPv6::Mapped.new(v.to_s) if v.is_a?(IPAddress::IPv4)
-    v = IPAddress(v.to_s) unless v.is_a?(IPAddress::IPv6)
+    if v.is_a?(Fixnum) || v.is_a?(Bignum)
+      v = if v < 2**32
+        IPAddress::IPv6::Mapped.parse_u128(v)
+      else
+        IPAddress::IPv6::parse_u128(v)
+      end
+    end
+      
 
     if v.respond_to?(:to_u128)
-      address = v.to_u128
+      self.address = v.to_u128
     else
       raise ArgumentError.new("Can't set #{v.class} to an IPv6Address")      
     end  
   end  
 end
 
-
-
 # Implementation of the BitCoin wire protocol, written using bindata.
 # Reference: https://en.bitcoin.it/wiki/Protocol_specification
 # 
diff --git a/spec/btc_wire_proto_spec.rb b/spec/btc_wire_proto_spec.rb
index 67fdf53..2785813 100644
--- a/spec/btc_wire_proto_spec.rb
+++ b/spec/btc_wire_proto_spec.rb
@@ -5,6 +5,62 @@ require 'btc_wire_proto'
 
 include ::BtcWireProto
 
+describe Ipv6Address do
+  describe "set" do
+    before(:each) { @ip = Ipv6Address.new }
+
+    it "should parse IPv4 strings into IPAddress::IPv6::Mapped objects" do
+      @ip.set("127.0.0.1")
+      @ip.get.should be_a(IPAddress::IPv6::Mapped)
+      @ip.get.to_s.should == "::ffff:127.0.0.1"
+    end
+
+    it "should parse IPv6 strings into IPAddress::IPv6 objects" do
+      @ip.set("fe80::1")
+      @ip.get.should be_a(IPAddress::IPv6)
+      @ip.get.to_s.should == "fe80::1"
+    end
+
+    it "should accept IPAddress::IPv(4|6) objects" do
+      @ip.set(IPAddress("127.0.0.1"))
+      @ip.get.should be_a(IPAddress::IPv6::Mapped)
+      @ip.get.to_s.should == "::ffff:127.0.0.1"
+
+      @ip.set(IPAddress("fe80::1"))
+      @ip.get.should be_a(IPAddress::IPv6)
+      @ip.get.to_s.should == "fe80::1"
+    end
+    
+    it "should accept arbitrary objects with #to_u128" do
+      o = Object.new
+      class << o
+        def to_u128 ; 666 ; end
+      end
+
+      @ip.set(o)
+      @ip.address.should == 666
+
+    end
+    
+    it "should treat Fixnums and Bignums as decimal IPs" do
+      ipv4 = IPAddress("127.0.0.1") # Should still be mapped
+      ipv6 = IPAddress("fe80::1")
+      
+      @ip.set(ipv4.to_u32)
+      @ip.get.should be_a(IPAddress::IPv6::Mapped)
+      @ip.get.to_s.should == "::ffff:127.0.0.1"
+      
+      @ip.set(ipv6.to_u128)
+      @ip.get.should be_a(IPAddress::IPv6)
+      @ip.get.to_s.should == "fe80::1"
+    end
+
+    it "should raise an ArgumentError when it can't parse objects" do
+      lambda { @ip.set("bad IP") }.should raise_error(ArgumentError)
+    end
+  end
+end
+
 describe ::BtcWireProto do
 
   # Payload fragments