Avoid calling count on Enumerable

lib/first_and_only.rb

@@ -2,15 +2,22 @@
 
 module Enumerable
   def first_and_only!
-    fail(FirstAndOnly::CountNotOne.new count) if first(2).count != 1
-    first
+    case first(2).count
+    when 0
+      fail FirstAndOnly::CountZero, FirstAndOnly::COUNT_ZERO_ERROR_MESSAGE
+    when 1
+      first
+    else
+      fail FirstAndOnly::CountGreaterThanOne, FirstAndOnly::COUNT_GREATER_THAN_ONE_ERROR_MESSAGE
+    end
   end
 
   module FirstAndOnly
-    class CountNotOne < StandardError
-      def initialize(count)
-        super("Expected the count to be 1, but it was #{count}.")
-      end
-    end
+    CountNotOne = Class.new(StandardError)
+    CountZero = Class.new(CountNotOne)
+    CountGreaterThanOne = Class.new(CountNotOne)
+
+    COUNT_ZERO_ERROR_MESSAGE = 'Expected the count to be 1, was 0.'.freeze
+    COUNT_GREATER_THAN_ONE_ERROR_MESSAGE = 'Expected the count to be 1, was greater than 1.'.freeze
   end
 end

spec/first_and_only_spec.rb

@@ -9,23 +9,33 @@
   shared_examples_for "it_is_empty" do
     specify do
       expect { subject.first_and_only! }.to raise_error(
-        Enumerable::FirstAndOnly::CountNotOne,
-        "Expected the count to be 1, but it was 0."
+        Enumerable::FirstAndOnly::CountZero,
+        "Expected the count to be 1, was 0."
       )
     end
   end
 
   shared_examples_for "it_has_more_than_one_element" do
     specify do
-      fail "bad example" if subject.count < 2
+      fail "bad example" if subject.first(2).count < 2
 
       expect { subject.first_and_only! }.to raise_error(
-        Enumerable::FirstAndOnly::CountNotOne,
-        "Expected the count to be 1, but it was #{subject.count}."
+        Enumerable::FirstAndOnly::CountGreaterThanOne,
+        "Expected the count to be 1, was greater than 1."
       )
     end
   end
 
+  describe "backwards compatibility with CountNotOne error class" do
+    specify do
+      expect(Enumerable::FirstAndOnly::CountZero).to be < Enumerable::FirstAndOnly::CountNotOne
+    end
+
+    specify do
+      expect(Enumerable::FirstAndOnly::CountGreaterThanOne).to be < Enumerable::FirstAndOnly::CountNotOne
+    end
+  end
+
   describe "an array" do
 
     context "with one element" do
@@ -76,4 +86,22 @@ def each
 
   end
 
+  describe "infinite enumerable" do
+    let(:infinite_enumerable) do
+      Class.new do
+        include Enumerable
+
+        def each
+          yield(rand(1_000_000)) while true
+        end
+      end
+    end
+
+    context "with more than on element" do
+      subject { infinite_enumerable.new }
+      it_behaves_like "it_has_more_than_one_element"
+    end
+
+  end
+
 end