import unittest
import psycopg2
import pandas as pd

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_two_ctes(self):
        """Check that query uses at least two CTEs."""
        with open("query.sql", "r") as file:
            query = file.read().lower()
        cte_count = query.count('with ')
        comma_cte = query.count('),')
        # Accepts either WITH ... , ... or WITH ... AS (...) , ... AS (...)
        condition = (cte_count == 1 and comma_cte >= 1) or (cte_count >= 2)
        _dynamic_test(
            self,
            condition,
            "Query uses at least two CTEs.",
            "Query must use at least two CTEs, with one nested within the other."
        )

    def test_only_high_salary_departments(self):
        """Check that only employees from departments with avg salary > 60000 are included."""
        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)
        # Compute actual qualifying departments
        self.cursor.execute("""
            SELECT d.department_id
            FROM departments d
            JOIN employees e ON d.department_id = e.department_id
            GROUP BY d.department_id
            HAVING AVG(e.salary) > 60000
        """)
        qualifying_ids = set(r[0] for r in self.cursor.fetchall())
        # Check all returned employees are from qualifying departments
        if not result_df.empty:
            self.cursor.execute("SELECT employee_id, department_id FROM employees")
            emp_df = pd.DataFrame(self.cursor.fetchall(), columns=["employee_id","department_id"])
            merged = pd.merge(result_df, emp_df, left_on="employee_id", right_on="employee_id", how="left")
            all_in = merged["department_id"].isin(qualifying_ids).all()
        else:
            all_in = True
        _dynamic_test(
            self,
            all_in,
            "Only employees from high average salary departments are included.",
            "Result includes employees from departments with average salary of $60,000 or less."
        )

    def test_result_columns(self):
        """Check that result includes correct columns."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        columns = [desc[0] for desc in self.cursor.description]
        expected = {"employee_id", "first_name", "last_name", "salary", "department_name"}
        has_all = expected.issubset(set(columns))
        _dynamic_test(
            self,
            has_all,
            "Result includes all required columns.",
            f"Result must include columns: {', '.join(expected)}."
        )

    def test_no_low_salary_departments(self):
        """Check that no employees from departments with avg salary <= 60000 are present."""
        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)
        # Find departments with avg salary <= 60000
        self.cursor.execute("""
            SELECT d.department_id
            FROM departments d
            JOIN employees e ON d.department_id = e.department_id
            GROUP BY d.department_id
            HAVING AVG(e.salary) <= 60000
        """)
        low_salary_ids = set(r[0] for r in self.cursor.fetchall())
        if not result_df.empty:
            self.cursor.execute("SELECT employee_id, department_id FROM employees")
            emp_df = pd.DataFrame(self.cursor.fetchall(), columns=["employee_id","department_id"])
            merged = pd.merge(result_df, emp_df, left_on="employee_id", right_on="employee_id", how="left")
            none_in = (~merged["department_id"].isin(low_salary_ids)).all()
        else:
            none_in = True
        _dynamic_test(
            self,
            none_in,
            "No employees from low average salary departments included.",
            "Result includes employees from departments with average salary of $60,000 or less."
        )

    def test_all_employees_in_qualifying_departments_listed(self):
        """Check that all employees from qualifying departments are listed."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Get qualifying department IDs
        self.cursor.execute("""
            SELECT d.department_id
            FROM departments d
            JOIN employees e ON d.department_id = e.department_id
            GROUP BY d.department_id
            HAVING AVG(e.salary) > 60000
        """)
        qualifying_ids = set(r[0] for r in self.cursor.fetchall())
        # Get all employees in those departments
        self.cursor.execute("SELECT employee_id FROM employees WHERE department_id = ANY(%s)", (list(qualifying_ids),))
        expected_emp_ids = set(r[0] for r in self.cursor.fetchall())
        # Get result employee ids
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(rows, columns=columns)
        result_emp_ids = set(result_df["employee_id"]) if not result_df.empty else set()
        missing = expected_emp_ids - result_emp_ids
        _dynamic_test(
            self,
            len(missing) == 0,
            "All employees in qualifying departments are listed.",
            f"Missing employees with IDs: {', '.join(str(i) for i in missing)}."
        )


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: Departmental Salary Insights

解答