class TestQueryBuilderAggregation(TestQueryBuilderBase):
    def setUp(self):
        self.builder = QueryBuilder()
        self.builder.add_table('table_a', 'a')
        self.builder.add_field('sum(a.f1)', 'fa')
        self.builder.add_field('a.f2', 'ga')

        self.builder.add_table('table_b', 'b')
        self.builder.add_field('count(b.f)', 'fb')
        self.builder.link_tables('a.id = b.id')

        self.builder.add_field('sum(a.f1) / count(b.f)', 'fab')
        self.builder.add_field('concat(a.f2, b.f)', 'cab')
        self.builder.add_field('0', 'c')

    def testSelectSumA(self):
        query = self.builder.select(['fa'])
        self.assertQuery(query, '''
            select sum(a.f1) fa
            from table_a a''')

    def testSelectCountB(self):
        query = self.builder.select(['fb'])
        self.assertQuery(query, '''
            select count(b.f) fb
            from table_b b''')

    def testSelectAvgAB(self):
        query = self.builder.select(['fab'])
        self.assertQuery(query, '''
            select sum(a.f1) / count(b.f) fab
            from table_a a,
                table_b b
            where a.id = b.id''')

    def testGroupBy(self):
        query = self.builder.group_by(['ga']).select(['ga', 'fa'])
        self.assertQuery(query, '''
            select a.f2 ga, sum(a.f1) fa
            from table_a a
            group by a.f2''')

        query = self.builder.group_by(['ga']).select(['fa'])
        self.assertQuery(query, '''
            select sum(a.f1) fa
            from table_a a
            group by a.f2''')

    def testGroupByFormula(self):
        query = self.builder.group_by(['cab']).select(['fa'])
        self.assertQuery(query, '''
            select sum(a.f1) fa
            from table_a a,
                table_b b
            where a.id = b.id
            group by concat(a.f2, b.f)''')

        query = self.builder.group_by(['ga', 'cab']).select(['fa'])
        self.assertQuery(query, '''
            select sum(a.f1) fa
            from table_a a,
                table_b b
            where a.id = b.id
            group by a.f2, concat(a.f2, b.f)''')

    def testHaving(self):
        query = self.builder.having('fa > 1').select(['fa'])
        self.assertQuery(query, '''
            select sum(a.f1) fa
            from table_a a
            having fa > 1''')

        query = self.builder.group_by(['ga']).having('fa > 1').select(['fa'])
        self.assertQuery(query, '''
            select sum(a.f1) fa
            from table_a a
            group by a.f2
            having fa > 1''')

    def testHavingContainUnselectedAlias(self):
        with self.assertRaisesRegexp(Exception, "alias 'cab' needs to be selected"):
            self.builder.having('cab = "abc"').select(['fa'])

    def testSelectConstant(self):
        query = self.builder.select(['c'])
        self.assertQuery(query, 'select 0 c')

        query = self.builder.select(['fa', 'c'])
        self.assertQuery(query, '''
            select sum(a.f1) fa, 0 c
            from table_a a''')
class TestQueryBuilder(TestQueryBuilderBase):
    def setUp(self):
        # table relationships
        self.builder = QueryBuilder()
        self.builder.add_table('table_a', 'a')
        self.builder.add_field('a.f', 'fa')

        self.builder.add_table('table_b', 'b')
        self.builder.link_tables('a.id = b.id')
        self.builder.add_field('b.f', 'fb')

        self.builder.add_table('table_c', 'c')
        self.builder.link_tables('a.id = c.id')
        self.builder.add_field('c.f', 'fc')

        self.builder.add_table('table_d', 'd')
        self.builder.link_tables('c.id = d.id')
        self.builder.add_field('d.f', 'fd')

    def testSelectA(self):
        query = self.builder.select(['fa'])
        self.assertQuery(query, 'select a.f fa from table_a a')

    def testSelectB(self):
        query = self.builder.select(['fb'])
        self.assertQuery(query, 'select b.f fb from table_b b')

    def testSelectC(self):
        query = self.builder.select(['fc'])
        self.assertQuery(query, 'select c.f fc from table_c c')

    def testSelectD(self):
        query = self.builder.select(['fd'])
        self.assertQuery(query, 'select d.f fd from table_d d')

    def testSelectAB(self):
        query = self.builder.select(['fa', 'fb'])
        self.assertQuery(query, '''
            select a.f fa, b.f fb 
            from table_a a,
                table_b b 
            where a.id = b.id''')

    def testSelectAC(self):
        query = self.builder.select(['fa', 'fc'])
        self.assertQuery(query, '''
            select a.f fa, c.f fc 
            from table_a a,
                table_c c 
            where a.id = c.id''')

    def testSelectAD(self):
        query = self.builder.select(['fa', 'fd'])
        self.assertQuery(query, '''
            select a.f fa, d.f fd 
            from table_a a,
                table_c c, 
                table_d d 
            where a.id = c.id and c.id = d.id''')

    def testSelectBC(self):
        query = self.builder.select(['fb', 'fc'])
        self.assertQuery(query, '''
            select b.f fb, c.f fc 
            from table_a a,
                table_b b,
                table_c c 
            where a.id = b.id and a.id = c.id''')

    def testSelectBD(self):
        query = self.builder.select(['fb', 'fd'])
        self.assertQuery(query, '''
            select b.f fb, d.f fd 
            from table_a a,
                table_b b, 
                table_c c, 
                table_d d 
            where a.id = b.id and a.id = c.id and c.id = d.id''')

    def testSelectCD(self):
        query = self.builder.select(['fc', 'fd'])
        self.assertQuery(query, '''
            select c.f fc, d.f fd 
            from table_c c,
                table_d d 
            where c.id = d.id''')

    def testSelectABC(self):
        query = self.builder.select(['fa', 'fb', 'fc'])
        self.assertQuery(query, '''
            select a.f fa, b.f fb, c.f fc 
            from table_a a,
                table_b b, 
                table_c c 
            where a.id = b.id and a.id = c.id''')

    def testSelectABCD(self):
        query = self.builder.select(['fa', 'fb', 'fc', 'fd'])
        self.assertQuery(query, '''
            select a.f fa, b.f fb, c.f fc, d.f fd 
            from table_a a, 
                table_b b, 
                table_c c, 
                table_d d
            where a.id = b.id and a.id = c.id and c.id = d.id''')

    def testSelectAWhereA(self):
        query = self.builder.where('fa = 1').select(['fa'])
        self.assertQuery(query, '''
            select a.f fa
            from table_a a
            where a.f = 1''')

        query = self.builder.where('fa = "a"').select(['fa'])
        self.assertQuery(query, '''
            select a.f fa
            from table_a a
            where a.f = "a"''')

        query = self.builder.where('fa = \'a\'').select(['fa'])
        self.assertQuery(query, '''
            select a.f fa
            from table_a a
            where a.f = \'a\'''')

    def testSelectAWhereB(self):
        query = self.builder.where('fb = 1').select(['fa'])
        self.assertQuery(query, '''
            select a.f fa
            from table_a a,
                table_b b
            where a.id = b.id and b.f = 1''')

    def testWhereChainable(self):
        query = self.builder.where('fa = 1').where('fb = 1').select(['fa'])
        self.assertQuery(query, '''
            select a.f fa
            from table_a a,
                table_b b
            where a.id = b.id and a.f = 1 and b.f = 1''')

    def testWhereReturnNewInstance(self):
        self.builder.where('fb = 1').select(['fa'])
        query = self.builder.where('fa = 1').select(['fa'])
        self.assertQuery(query, '''
            select a.f fa
            from table_a a
            where a.f = 1''')