Return `nil` if postcodes don't match exactly.

2 files changed, 38 insertions, 13 deletions

diff --git a/lib/norwegian-postcodes.rb b/lib/norwegian-postcodes.rb
index cc3cd31..be77c8e 100644
--- a/lib/norwegian-postcodes.rb
+++ b/lib/norwegian-postcodes.rb
@@ -40,26 +40,31 @@ module PostCodes
       [code, PostCodes.county(code)]
     end
 
-    def >=(postcode)
-      @postcode.to_i >= postcode.to_i
-    end
-
     def to_s
       [@postcode, @city, @municipality, @municipality_name, @cat].join("\t")
     end
   end
 
   class << self
-    def load(f)
+    def load(file)
       @postcodes = []
-      IO.foreach(f, :encoding => Encoding::ISO_8859_15) do |l|
+      if file.is_a?(String)
+        f = File.open(file, :encoding => Encoding::ISO_8859_15)
+      else
+        f = file
+      end
+
+      f.each_line do |l|
         a = l.chomp().split("\t").map{|s| s.encode(Encoding::UTF_8)}
         @postcodes << PostCode.new(*a)
       end
     end
 
     def search(pc)
-      @postcodes.bsearch {|x| x >= pc}
+      res = @postcodes.bsearch {|x| x.postcode.to_i >= pc.to_i}
+      unless res.nil?
+        res.postcode.to_i == pc.to_i ? res : nil
+      end
     end
 
     def county(c)
diff --git a/test/test_postcodes.rb b/test/test_postcodes.rb
index e3f2080..3b38784 100644
--- a/test/test_postcodes.rb
+++ b/test/test_postcodes.rb
@@ -6,12 +6,6 @@ describe PostCodes::PostCode do
     @postcode = PostCodes::PostCode.new('1234', 'MYTOWN', '1378', 'MY MUNICIPALITY', 'G')
   end
 
-  it "compares to a postcode" do
-    (@postcode >= '1233').must_equal true
-    (@postcode >= '1234').must_equal true
-    (@postcode >= '1235').must_equal false
-  end
-
   it "has a county" do
     @postcode.county.must_equal [13, '(BERGEN)']
   end
@@ -53,3 +47,29 @@ describe 'PostCodes::county' do
     PostCodes.county(24).must_equal nil
   end
 end
+
+describe PostCodes do
+  before do
+    codes = [
+      "0100\tFIRST TOWN\t0101\tFIRST MUNICIPALITY\tG",
+      "1234\tMYTONW\t1278\tMY MUNICIPALITY\tG",
+      "6734\tOTHERTOWN\t2783\tOTHER MUNICIPALITY\tG"
+    ]
+    PostCodes.load(StringIO.new(codes.join("\n").encode(Encoding::ISO_8859_15)))
+  end
+
+  it "can find existing postcode" do
+    c = PostCodes.search('0100')
+    c.must_be_kind_of(PostCodes::PostCode)
+    c.city.must_equal('FIRST TOWN')
+
+    c1 = PostCodes.search('6734')
+    c1.must_be_kind_of(PostCodes::PostCode)
+    c1.city.must_equal('OTHERTOWN')
+  end
+
+  it "returns nil if postcode does not exist" do
+    c = PostCodes.search('2222')
+    c.must_equal(nil)
+  end
+end