# -*- coding: utf-8 -*- from __future__ import absolute_import from __future__ import division from __future__ import print_function from __future__ import unicode_literals from superset.dataframe import dedup, SupersetDataFrame from superset.db_engine_specs import BaseEngineSpec from .base_tests import SupersetTestCase class SupersetDataFrameTestCase(SupersetTestCase): def test_dedup(self): self.assertEquals( dedup(['foo', 'bar']), ['foo', 'bar'], ) self.assertEquals( dedup(['foo', 'bar', 'foo', 'bar']), ['foo', 'bar', 'foo__1', 'bar__1'], ) self.assertEquals( dedup(['foo', 'bar', 'bar', 'bar']), ['foo', 'bar', 'bar__1', 'bar__2'], ) def test_get_columns_basic(self): data = [ ('a1', 'b1', 'c1'), ('a2', 'b2', 'c2'), ] cursor_descr = ( ('a', 'string'), ('b', 'string'), ('c', 'string'), ) cdf = SupersetDataFrame(data, cursor_descr, BaseEngineSpec) self.assertEqual( cdf.columns, [ { 'is_date': False, 'type': 'STRING', 'name': 'a', 'is_dim': True, }, { 'is_date': False, 'type': 'STRING', 'name': 'b', 'is_dim': True, }, { 'is_date': False, 'type': 'STRING', 'name': 'c', 'is_dim': True, }, ], ) def test_get_columns_with_int(self): data = [ ('a1', 1), ('a2', 2), ] cursor_descr = ( ('a', 'string'), ('b', 'int'), ) cdf = SupersetDataFrame(data, cursor_descr, BaseEngineSpec) self.assertEqual( cdf.columns, [ { 'is_date': False, 'type': 'STRING', 'name': 'a', 'is_dim': True, }, { 'is_date': False, 'type': 'INT', 'name': 'b', 'is_dim': False, 'agg': 'sum', }, ], ) def test_get_columns_type_inference(self): data = [ (1.2, 1), (3.14, 2), ] cursor_descr = ( ('a', None), ('b', None), ) cdf = SupersetDataFrame(data, cursor_descr, BaseEngineSpec) self.assertEqual( cdf.columns, [ { 'is_date': False, 'type': 'FLOAT', 'name': 'a', 'is_dim': False, 'agg': 'sum', }, { 'is_date': False, 'type': 'INT', 'name': 'b', 'is_dim': False, 'agg': 'sum', }, ], )