collapse a couple of methods using find_all.main

2 files changed, 10 insertions, 11 deletions

diff --git a/lib/grant.rb b/lib/grant.rb
index a993366..10f89f4 100644
--- a/lib/grant.rb
+++ b/lib/grant.rb
@@ -1,4 +1,5 @@
 class Grant
+  include Enumerable
   attr_reader :issued_at
 
   def initialize(value_of_grant, share_price, date = Clock.today)
@@ -7,14 +8,6 @@ class Grant
     @issued_at = date
   end
 
-  def value_of_unvested_units_at(price)
-    value_of(unvested_units, price)
-  end
-
-  def value_of_vested_units_at(price)
-    value_of(vested_units, price)
-  end
-
   def vest_at(price, portion)
     number_of_units_to_vest = portion.to_f * @units.count
     unvested_units.take(number_of_units_to_vest).each do |unit|
@@ -22,7 +15,9 @@ class Grant
     end
   end
 
-  private
+  def each
+    @units.each { |unit| yield unit }
+  end
 
   def value_of(units, price)
     units.inject(0.00.dollars) do |memo, unit|
@@ -30,6 +25,8 @@ class Grant
     end
   end
 
+  private
+
   def unvested_units
     @units - vested_units
   end
diff --git a/lib/unit.rb b/lib/unit.rb
index 32c9a4d..16e6456 100644
--- a/lib/unit.rb
+++ b/lib/unit.rb
@@ -19,12 +19,14 @@ end
 
 class UnvestedUnits
   def from(grant, at_price)
-    grant.value_of_unvested_units_at(at_price)
+    unvested_units = grant.find_all { |unit| unit.vested? == false }
+    grant.value_of(unvested_units, at_price)
   end
 end
 
 class VestedUnits
   def from(grant, at_price)
-    grant.value_of_vested_units_at(at_price)
+    vested_units = grant.find_all { |unit| unit.vested? }
+    grant.value_of(vested_units, at_price)
   end
 end