import unittest
import psycopg2
import pandas as pd
import re

class TestPostgresQuery(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.connection = psycopg2.connect(
            host="localhost",
            user="agent",
            password="qwerty",
            dbname="testdb",
            sslmode="disable"
        )
        cls.cursor = cls.connection.cursor()

    def test_uses_cte(self):
        """Check that a CTE is used in the query."""
        with open("query.sql", "r") as file:
            query = file.read()
        condition = normalize_text("with ") in normalize_text(query)
        _dynamic_test(
            self,
            condition,
            "Query uses a CTE.",
            "Query does not use a CTE."
        )

    def test_joins_departments_employees_sales(self):
        """Check that the query joins departments, employees, and sales tables."""
        with open("query.sql", "r") as file:
            query = file.read()
        norm_query = normalize_text(query)
        has_departments = "departments" in norm_query
        has_employees = "employees" in norm_query
        has_sales = "sales" in norm_query
        condition = has_departments and has_employees and has_sales
        _dynamic_test(
            self,
            condition,
            "Query joins all required tables.",
            "Query does not join all required tables."
        )

    def test_selects_only_departments_with_sales_gt_10000(self):
        """Check that only departments with total sales > 10000 are returned."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(rows, columns=columns)
        # All total_sales must be > 10000
        if not result_df.empty and 'total_sales' in result_df.columns:
            condition = (result_df['total_sales'] > 10000).all()
        else:
            condition = True  # If no rows returned, that's acceptable
        _dynamic_test(
            self,
            condition,
            "All returned departments have total_sales greater than 10000.",
            "Some returned departments have total_sales less than or equal to 10000."
        )

    def test_output_fields(self):
        """Check that output includes department_id, department_name, and total_sales fields."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        columns = [desc[0] for desc in self.cursor.description]
        expected_cols = {'department_id', 'department_name', 'total_sales'}
        condition = expected_cols.issubset(set(columns))
        _dynamic_test(
            self,
            condition,
            "Output contains all required columns.",
            "Output is missing one or more required columns."
        )

    def test_correct_aggregation(self):
        """Check that the aggregation is correct for departments with sales > 10000."""
        # Compute expected department sales
        solution_query = """
        WITH department_sales AS (
          SELECT
            d.department_id,
            d.department_name,
            SUM(s.amount) AS total_sales
          FROM
            departments d
            JOIN employees e ON d.department_id = e.department_id
            JOIN sales s ON e.employee_id = s.employee_id
          GROUP BY
            d.department_id,
            d.department_name
        )
        SELECT
          department_id,
          department_name,
          total_sales
        FROM
          department_sales
        WHERE
          total_sales > 10000;
        """
        self.cursor.execute(solution_query)
        expected_rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        expected_df = pd.DataFrame(expected_rows, columns=columns)
        # Now run user query
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        user_rows = self.cursor.fetchall()
        user_columns = [desc[0] for desc in self.cursor.description]
        user_df = pd.DataFrame(user_rows, columns=user_columns)
        # Compare dataframes for matching rows
        missing_rows = []
        for _, expected_row in expected_df.iterrows():
            expected_dict = expected_row.to_dict()
            mask = (user_df[list(expected_dict.keys())] == pd.Series(expected_dict)).all(axis=1)
            if not mask.any():
                missing_rows.append(expected_dict)
        condition = len(missing_rows) == 0
        if not condition:
            human_readable = "\\n".join(str(r) for r in missing_rows)
            failure_msg = f"The following expected rows are missing:\\n{human_readable}"
        else:
            failure_msg = ""
        _dynamic_test(
            self,
            condition,
            "Query returns correct aggregated results.",
            failure_msg
        )

def _dynamic_test(test_case, condition, success_message, failure_message):
    if condition:
        test_case._testMethodName = success_message
        test_case.assertTrue(True, success_message)
    else:
        test_case._testMethodName = failure_message
        test_case.fail(failure_message)

def normalize_text(text):
    text = text.lower()
    text = re.sub(r"\s{2,}", " ", text)
    text = re.sub(r"\s*([,:?])\s*", r"\1 ", text)
    return text.strip()

if __name__ == "__main__":
    unittest.main()


test_solution.py

A comprehensive course designed to help you master Common Table Expressions (CTEs) in SQL, from foundational concepts to advanced real-world applications. Explore recursive queries, modular query design, and performance considerations through engaging theory and practical, scenario-based tasks.

Explore the basics of Common Table Expressions, their syntax, and how they can simplify SQL queries. Practice using CTEs in real-world scenarios such as reporting and data transformation.

Delve into recursive CTEs for handling hierarchical data, such as organizational charts and bill of materials. Practice building recursive queries to solve real-world problems.

Master advanced CTE techniques, including combining CTEs with aggregates, optimizing performance, and using CTEs for complex reporting.

Challenge: Department Sales Aggregation

Solução