improve the partition_by_shard stub

test plan: n/a

Change-Id: I35abecfc43d776f2d21c63569fe197b6e461dfe8
Reviewed-on: https://gerrit.instructure.com/8841
Tested-by: Hudson <hudson@instructure.com>
Reviewed-by: Bracken Mosbacker <bracken@instructure.com>

lib/stubs/shard.rb

@@ -29,8 +29,8 @@ class Shard
     default
   end
 
-  def self.partition_by_shard(array)
-    yield array
+  def self.partition_by_shard(array, partition_proc = nil)
+    Array(yield array)
   end
 
   def activate