Add tests

superset/models/core.py

@@ -786,7 +786,12 @@ def grains(self):
         return self.db_engine_spec.time_grains
 
     def grains_dict(self):
-        return {grain.duration: grain for grain in self.grains()}
+        """Allowing to lookup grain by either label or duration
+
+        For backward compatibility"""
+        d = {grain.duration: grain for grain in self.grains()}
+        d.update({grain.label: grain for grain in self.grains()})
+        return d
 
     def get_extra(self):
         extra = {}

tests/model_tests.py

@@ -97,3 +97,11 @@ def test_select_star(self):
         FROM bart_lines
         LIMIT 100""".format(**locals()))
         assert sql.startswith(expected)
+
+    def test_grains_dict(self):
+        uri = 'mysql://root@localhost'
+        database = Database(sqlalchemy_uri=uri)
+        d = database.grains_dict()
+        self.assertEquals(d.get('day').function, 'DATE({col})')
+        self.assertEquals(d.get('P1D').function, 'DATE({col})')
+        self.assertEquals(d.get('Time Column').function, '{col}')