From ef8b8156babafcdea73784900296ea618b0d10dd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Milka?= <skubekk@gmail.com>
Date: Thu, 31 May 2018 14:36:33 +0200
Subject: [PATCH] feat(taskscheduler): allow to use custom scheduler for shards

When shardTestFiles is true, specs will be sharded by file. (i.e. every test file is executed in separated browser instance).
This behavior can now be customized by providing custom shard scheduler function which allows to run multiple specs in one browser instance. For example we can run same amount of specs in each browser instance.
---
 lib/config.ts                   |  7 ++++-
 lib/taskScheduler.ts            | 22 ++++++++-------
 spec/unit/taskScheduler_test.js | 48 +++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 11 deletions(-)

diff --git a/lib/config.ts b/lib/config.ts
index 043ed9202..c75cca62b 100644
--- a/lib/config.ts
+++ b/lib/config.ts
@@ -1,5 +1,7 @@
 import {PluginConfig} from './plugins';
 
+export type SpecScheduler = () => (specs: Array<string>, capabilities: any) => Array<Array<string>>;
+
 export interface Config {
   [key: string]: any;
 
@@ -325,9 +327,12 @@ export interface Config {
     /**
      * If this is set to be true, specs will be sharded by file (i.e. all
      * files to be run by this set of capabilities will run in parallel).
+     * By default each test file will run in separeted browser instance.
      * Default is false.
+     * You can also provide custom spec scheduler function
+     * that can change default behaviour and run many specs in one browser instance.
      */
-    shardTestFiles?: boolean;
+    shardTestFiles: boolean | SpecScheduler;
 
     /**
      * Maximum number of browser instances that can run in parallel for this
diff --git a/lib/taskScheduler.ts b/lib/taskScheduler.ts
index 420c51dd0..de9a56aeb 100644
--- a/lib/taskScheduler.ts
+++ b/lib/taskScheduler.ts
@@ -63,17 +63,19 @@ export class TaskScheduler {
         });
       }
 
-      let specLists: Array<Array<string>> = [];
-      // If we shard, we return an array of one element arrays, each containing
-      // the spec file. If we don't shard, we return an one element array
-      // containing an array of all the spec files
-      if (capabilities.shardTestFiles) {
-        capabilitiesSpecs.forEach((spec) => {
-          specLists.push([spec]);
-        });
-      } else {
-        specLists.push(capabilitiesSpecs);
+      let shardScheduler = capabilities.shardTestFiles;
+      if (typeof shardScheduler !== 'function') {
+        // If we shard, we return an array of one element arrays, each containing
+        // the spec file. If we don't shard, we return an one element array
+        // containing an array of all the spec files
+        shardScheduler = function(specs: Array<string>, capabilities: any): Array<Array<string>> {
+          if (capabilities.shardTestFiles) {
+            return specs.map(spec => [spec]);
+          }
+          return [specs];
+        };
       }
+      const specLists = shardScheduler(capabilitiesSpecs, capabilities);
 
       capabilities.count = capabilities.count || 1;
 
diff --git a/spec/unit/taskScheduler_test.js b/spec/unit/taskScheduler_test.js
index 8bc19a053..f8e195fa8 100644
--- a/spec/unit/taskScheduler_test.js
+++ b/spec/unit/taskScheduler_test.js
@@ -258,4 +258,52 @@ describe('the task scheduler', function() {
     expect(scheduler.numTasksOutstanding()).toEqual(0);
   });
 
+  it('should use custom shard scheduler when provided', function() {
+    var toAdd = {
+      specs: [
+        'spec/unit/data/fakespecA.js',
+        'spec/unit/data/fakespecB.js',
+        'spec/unit/data/fakespecC.js'
+      ],
+      multiCapabilities: [{
+        shardTestFiles: shardScheduler,
+        browserName: 'chrome',
+        maxInstances: 2,
+      }]
+    };
+    var config = new ConfigParser().addConfig(toAdd).getConfig();
+    var scheduler = new TaskScheduler(config);
+
+    var task1 = scheduler.nextTask();
+    expect(task1.capabilities.browserName).toEqual('chrome');
+    expect(task1.specs.length).toEqual(2);
+
+    var task2 = scheduler.nextTask();
+    expect(task2.capabilities.browserName).toEqual('chrome');
+    expect(task2.specs.length).toEqual(1);
+
+    task1.done();
+    task2.done();
+    expect(scheduler.numTasksOutstanding()).toEqual(0);
+  });
+
+  function shardScheduler(specs, capabilities) {
+    const numberOfShards = capabilities.maxInstances;
+    if(numberOfShards > 1) {
+        const bucketSize = Math.ceil(specs.length/numberOfShards);
+        const shards = [];
+        let start = 0;
+        while (start < specs.length) {
+            let end = start + bucketSize;
+            if ( end > specs.length) {
+                end = specs.length;
+            }
+            shards.push(specs.slice(start,end));
+            start = end;
+        }
+        return shards;
+    }
+    return [specs];
+  }
+
 });