diff --git a/src/utils/ColorGenerator.js b/src/utils/ColorGenerator.js
index 840b485b..d7a4f9e2 100644
--- a/src/utils/ColorGenerator.js
+++ b/src/utils/ColorGenerator.js
@@ -1,37 +1,46 @@
 import Storage from './Storage';
 import PropTypes from 'prop-types';
+import { range } from 'ramda';
 
-const buildRandomColor = () => {
-  const letters = '0123456789ABCDEF';
-  let color = '#';
-  for (let i = 0; i < 6; i++ ) {
-    color += letters[Math.floor(Math.random() * 16)];
-  }
-  return color;
-};
+const { floor, random } = Math;
+const letters = '0123456789ABCDEF';
+const buildRandomColor = () =>
+  `#${
+    range(0, 6)
+      .map(() => letters[floor(random() * 16)])
+      .join('')
+  }`;
 
 export class ColorGenerator {
   constructor(storage) {
     this.storage = storage;
     this.colors = this.storage.get('colors') || {};
+
+    this.getColorForKey = this.getColorForKey.bind(this);
+    this.setColorForKey = this.setColorForKey.bind(this);
   }
 
   getColorForKey = key => {
-    let color = this.colors[key];
-    if (color) {
-      return color;
-    }
+    const color = this.colors[key];
 
     // If a color has not been set yet, generate a random one and save it
-    color = buildRandomColor();
-    this.colors[key] = color;
-    this.storage.set('colors', this.colors);
+    if (!color) {
+      this.setColorForKey(key, buildRandomColor());
+      return this.getColorForKey(key);
+    }
+
     return color;
   };
+
+  setColorForKey = (key, color) => {
+    this.colors[key] = color;
+    this.storage.set('colors', this.colors);
+  }
 }
 
 export const colorGeneratorType = PropTypes.shape({
   getColorForKey: PropTypes.func,
+  setColorForKey: PropTypes.func,
 });
 
 export default new ColorGenerator(Storage);