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_salary_rank_top3(self):
        """Check only employees with salary_rank 1, 2, or 3 per department are returned."""
        with open("query.sql", "r") as file:
            query = file.read()

        # Execute user query
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(rows, columns=columns)

        # Calculate expected ranks for each department
        self.cursor.execute("SELECT employee_id, department, salary FROM employees")
        all_emps = pd.DataFrame(self.cursor.fetchall(), columns=["employee_id", "department", "salary"])
        all_emps["salary"] = all_emps["salary"].astype(float)
        all_emps["salary_rank"] = all_emps.groupby("department")['salary'].rank(method="min", ascending=False)
        top3 = all_emps[all_emps["salary_rank"] <= 3]

        # All employee_ids in result must be in top3 per department
        missing = []
        for _, row in result_df.iterrows():
            cond = (
                (top3["employee_id"] == row["employee_id"]) &
                (top3["department"] == row["department"])
            )
            if not cond.any():
                missing.append(dict(row))
        condition = len(missing) == 0
        fail_msg = f"Returned employees not in top 3 salaries per department: {missing}" if missing else ""
        _dynamic_test(self, condition, "All returned employees are in top 3 per department.", fail_msg)

    def test_all_departments_included(self):
        """Check all departments are represented, even if fewer than 3 employees."""
        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)
        self.cursor.execute("SELECT DISTINCT department FROM employees")
        all_depts = set(r[0] for r in self.cursor.fetchall())
        result_depts = set(result_df['department'].unique())
        missing = all_depts - result_depts
        condition = len(missing) == 0
        fail_msg = f"Missing departments in result: {missing}" if missing else ""
        _dynamic_test(self, condition, "All departments are present in the output.", fail_msg)

    def test_output_columns(self):
        """Check output columns are exactly employee_id, name, department, salary."""
        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 = ["employee_id", "name", "department", "salary"]
        condition = columns == expected_cols
        fail_msg = f"Output columns are {columns}, expected {expected_cols}."
        _dynamic_test(self, condition, "Output columns are correct.", fail_msg)

    def test_output_ordering(self):
        """Check results are ordered by department, salary descending, name."""
        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)
        # Create a sorted version
        sorted_df = result_df.sort_values(["department", "salary", "name"], ascending=[True, False, True]).reset_index(drop=True)
        result_df = result_df.reset_index(drop=True)
        condition = result_df.equals(sorted_df)
        fail_msg = "Rows are not ordered by department, salary descending, then name."
        _dynamic_test(self, condition, "Rows are ordered correctly.", 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

Master the art of analyzing data with SQL window functions. This course guides you through the core concepts, syntax, and practical applications of window functions, empowering you to perform advanced analytics and reporting directly in SQL.

Explore the foundational concepts of window functions, their syntax, and how they differ from traditional aggregate functions. Learn to apply basic window functions for analytics.

Dive into ranking, percentiles, and analytic functions. Learn to use ROW_NUMBER, RANK, DENSE_RANK, and NTILE for advanced reporting.

Apply window functions to solve complex business problems, including advanced framing, multi-function queries, and real-world analytics.

Challenge: Top 3 Salaries per Department

解答