collapse a couple of methods using find_all.main

1 files changed, 4 insertions, 2 deletions

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