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_query_uses_two_ctes(self):
        """Check that the query uses two CTEs."""
        with open("query.sql", "r") as file:
            query = file.read()
        cte_count = query.lower().count('with ')
        recursive_count = query.lower().count('with recursive')
        # Count commas after first WITH to count chained CTEs
        if cte_count > 0:
            cte_names = query.lower().split('with', 1)[1]
            cte_num = cte_names.count(' as (')
        else:
            cte_num = 0
        _dynamic_test(
            self,
            cte_num >= 2,
            "Query uses at least two CTEs as required.",
            "Query must use two CTEs as described."
        )

    def test_first_cte_employee_sales(self):
        """Check that the first CTE calculates total sales per employee."""
        # Reference solution result for employee_sales CTE
        solution = '''
            SELECT
                e.employee_id,
                e.first_name,
                e.last_name,
                e.department_id,
                SUM(s.amount) AS total_sales
            FROM
                employees e
                LEFT JOIN sales s ON e.employee_id = s.employee_id
            GROUP BY
                e.employee_id, e.first_name, e.last_name, e.department_id
        '''
        self.cursor.execute(solution)
        expected_rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        expected_df = pd.DataFrame(expected_rows, columns=columns)

        with open("query.sql", "r") as file:
            user_query = file.read()
        # Try to extract the employee_sales CTE and run it
        try:
            # Find the CTE definition
            lower_query = user_query.lower()
            start = lower_query.find('employee_sales as')
            if start == -1:
                _dynamic_test(
                    self,
                    False,
                    "employee_sales CTE calculates total sales per employee.",
                    "First CTE named employee_sales not found."
                )
                return
            # Find the opening and closing parenthesis for CTE
            open_paren = user_query.find('(', start)
            close_paren = open_paren
            depth = 1
            while depth > 0 and close_paren < len(user_query) - 1:
                close_paren += 1
                if user_query[close_paren] == '(': depth += 1
                elif user_query[close_paren] == ')': depth -= 1
            cte_sql = user_query[open_paren + 1:close_paren]
            cte_sql = cte_sql.strip()
            # Wrap CTE SQL in SELECT * FROM (...)
            wrapped_cte_sql = f"SELECT * FROM ({cte_sql}) AS employee_sales"
            self.cursor.execute(wrapped_cte_sql)
            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
            # Allow for NULL/None comparison in total_sales
            missing = []
            for _, row in expected_df.iterrows():
                expected_dict = row.to_dict()
                # For each expected row, check if a matching row exists in user_df
                found = False
                for _, user_row in user_df.iterrows():
                    match = True
                    for key in expected_dict:
                        if pd.isna(expected_dict[key]) and pd.isna(user_row[key]):
                            continue
                        if expected_dict[key] != user_row[key]:
                            match = False
                            break
                    if match:
                        found = True
                        break
                if not found:
                    missing.append(expected_dict)
            condition = len(missing) == 0
            if not condition:
                fail_msg = f"Your first CTE does not calculate total sales per employee correctly. Missing: {missing}"
            else:
                fail_msg = ""
            _dynamic_test(
                self,
                condition,
                "First CTE calculates total sales per employee correctly.",
                fail_msg
            )
        except Exception as e:
            _dynamic_test(
                self,
                False,
                "First CTE calculates total sales per employee.",
                f"Error executing/extracting first CTE: {e}"
            )

    def test_second_cte_department_avg_sales(self):
        """Check that the second CTE calculates average sales per department and selects employees whose total sales exceed the departmental average."""
        # Reference solution for final output
        solution = '''
        WITH employee_sales AS (
            SELECT
                e.employee_id,
                e.first_name,
                e.last_name,
                e.department_id,
                SUM(s.amount) AS total_sales
            FROM
                employees e
                LEFT JOIN sales s ON e.employee_id = s.employee_id
            GROUP BY
                e.employee_id, e.first_name, e.last_name, e.department_id
        ),
        department_avg_sales AS (
            SELECT
                es.employee_id,
                es.first_name,
                es.last_name,
                es.department_id,
                es.total_sales,
                AVG(es.total_sales) OVER (PARTITION BY es.department_id) AS dept_avg_sales
            FROM
                employee_sales es
        )
        SELECT
            employee_id,
            first_name,
            last_name,
            department_id,
            total_sales
        FROM
            department_avg_sales
        WHERE
            total_sales > dept_avg_sales;
        '''
        self.cursor.execute(solution)
        expected_rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        expected_df = pd.DataFrame(expected_rows, columns=columns)

        with open("query.sql", "r") as file:
            user_query = file.read()
        self.cursor.execute(user_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)

        missing = []
        for _, row in expected_df.iterrows():
            mask = (user_df[list(row.keys())] == pd.Series(row)).all(axis=1)
            if not mask.any():
                missing.append(row.to_dict())
        condition = len(missing) == 0
        if not condition:
            fail_msg = f"The following expected employees are missing from your result: {missing}"
        else:
            fail_msg = ""
        _dynamic_test(
            self,
            condition,
            "Second CTE selects employees exceeding department average sales correctly.",
            fail_msg
        )

    def test_result_columns(self):
        """Check that the result includes the correct columns."""
        with open("query.sql", "r") as file:
            user_query = file.read()
        self.cursor.execute(user_query)
        columns = [desc[0] for desc in self.cursor.description]
        expected_cols = ["employee_id", "first_name", "last_name", "department_id", "total_sales"]
        missing = [col for col in expected_cols if col not in columns]
        extra = [col for col in columns if col not in expected_cols]
        condition = len(missing) == 0 and len(extra) == 0
        if not condition:
            fail_msg = f"Your result columns are incorrect. Missing: {missing}, Extra: {extra}"
        else:
            fail_msg = ""
        _dynamic_test(
            self,
            condition,
            "Result includes required columns only.",
            fail_msg
        )

    def test_result_only_employees_above_department_average(self):
        """Check that only employees with total sales greater than their department's average are present in the result."""
        # Get department averages
        solution = '''
            WITH employee_sales AS (
                SELECT
                    e.employee_id,
                    e.first_name,
                    e.last_name,
                    e.department_id,
                    SUM(s.amount) AS total_sales
                FROM
                    employees e
                    LEFT JOIN sales s ON e.employee_id = s.employee_id
                GROUP BY
                    e.employee_id, e.first_name, e.last_name, e.department_id
            )
            SELECT
                es.employee_id,
                es.department_id,
                es.total_sales,
                AVG(es.total_sales) OVER (PARTITION BY es.department_id) AS dept_avg_sales
            FROM employee_sales es
        '''
        self.cursor.execute(solution)
        dept_df = pd.DataFrame(self.cursor.fetchall(), columns=[desc[0] for desc in self.cursor.description])
        dept_avg = dept_df.groupby('department_id')['dept_avg_sales'].first().to_dict()
        emp_sales = dept_df.set_index('employee_id')['total_sales'].to_dict()
        emp_dept = dept_df.set_index('employee_id')['department_id'].to_dict()
        # Get user result
        with open("query.sql", "r") as file:
            user_query = file.read()
        self.cursor.execute(user_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)
        # Check all employees in result
        failed = []
        for _, row in user_df.iterrows():
            eid = row['employee_id']
            dept = emp_dept.get(eid)
            sales = emp_sales.get(eid)
            avg = dept_avg.get(dept)
            if avg is None or sales is None or not (sales > avg):
                failed.append(eid)
        condition = len(failed) == 0
        if not condition:
            fail_msg = f"The following employee_ids should not be in the result: {failed}"
        else:
            fail_msg = ""
        _dynamic_test(
            self,
            condition,
            "Result includes only employees above department average.",
            fail_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: Multi-Step Employee Analysis

解答