import unittest
import psycopg2
import pandas as pd
import math

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_columns_present(self):
        """Check that required columns are present in the result."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        columns = [desc[0] for desc in self.cursor.description]
        expected_columns = set(["employee_id", "name", "avg_score", "performance_quartile"])
        result = set(columns)
        _dynamic_test(
            self,
            expected_columns.issubset(result),
            "All required columns are present.",
            f"Missing columns: {expected_columns - result}"
        )

    def test_all_employees_with_reviews_included(self):
        """Check that all employees with reviews are included in the result."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute("SELECT DISTINCT employee_id FROM performance")
        expected_ids = set(row[0] for row in self.cursor.fetchall())
        self.cursor.execute(query)
        result_ids = set(row[0] for row in self.cursor.fetchall())
        _dynamic_test(
            self,
            expected_ids == result_ids,
            "All employees with reviews are included.",
            f"Missing employee_ids: {expected_ids - result_ids}"
        )

    def test_avg_score_calculation(self):
        """Check that avg_score is calculated correctly for each employee."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Calculate expected averages
        self.cursor.execute("SELECT employee_id, AVG(score) FROM performance GROUP BY employee_id")
        expected_averages = {row[0]: float(row[1]) for row in self.cursor.fetchall()}
        self.cursor.execute(query)
        result = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        idx_id = columns.index("employee_id")
        idx_avg = columns.index("avg_score")
        mismatches = []
        for row in result:
            emp_id = row[idx_id]
            avg_score = float(row[idx_avg])
            if round(avg_score, 5) != round(expected_averages[emp_id], 5):
                mismatches.append((emp_id, avg_score, expected_averages[emp_id]))
        _dynamic_test(
            self,
            not mismatches,
            "Average scores are calculated correctly.",
            f"Incorrect averages for employee_ids: {mismatches}"
        )

    def test_ntile_quartile_assignment(self):
        """Check that NTILE(4) correctly assigns quartiles with 1 = highest avg_score."""
        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]
        df = pd.DataFrame(rows, columns=columns)
        df = df.sort_values("avg_score", ascending=False).reset_index(drop=True)
        n = len(df)
        # Use the same logic as Postgres NTILE to assign expected quartiles
        # NTILE(4): bucket_size = ceil(n/4), assign buckets in round-robin order
        bucket_size = math.ceil(n / 4)
        expected_quartiles = [((i // bucket_size) + 1) for i in range(n)]
        # If n is not divisible by 4, NTILE distributes extra rows to lower-numbered buckets
        # Let's use the logic from Postgres docs
        # See: https://www.postgresql.org/docs/current/functions-window.html#FUNCTIONS-WINDOW-TABLE
        # The NTILE function divides the ordered rows as evenly as possible among the number of buckets.
        # If the number of rows is not divisible by the number of buckets, the first buckets get one extra row each.
        ntiles = 4
        per_bucket = n // ntiles
        remainder = n % ntiles
        bucket_limits = []
        current = 0
        for i in range(1, ntiles+1):
            add = per_bucket + (1 if i <= remainder else 0)
            bucket_limits.append(current + add)
            current += add
        expected_quartiles = []
        for i in range(n):
            for b, lim in enumerate(bucket_limits):
                if i < lim:
                    expected_quartiles.append(b+1)
                    break
        df["expected_quartile"] = expected_quartiles
        mismatches = df[df["performance_quartile"] != df["expected_quartile"]]
        _dynamic_test(
            self,
            mismatches.empty,
            "NTILE(4) assigns quartiles correctly (1 = highest avg_score).",
            f"Quartile assignment mismatches: {mismatches[['employee_id','avg_score','performance_quartile','expected_quartile']].to_dict('records')}"
        )

    def test_quartile_1_is_highest(self):
        """Check that employees with highest avg_score are assigned to quartile 1."""
        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]
        df = pd.DataFrame(rows, columns=columns)
        if df.empty:
            _dynamic_test(self, False, "", "Query result is empty.")
            return
        max_avg = df["avg_score"].max()
        top_quartile = df[df["avg_score"] == max_avg]["performance_quartile"].unique()
        _dynamic_test(
            self,
            all(q == 1 for q in top_quartile),
            "Employees with highest avg_score are assigned to quartile 1.",
            f"Top scoring employees not in quartile 1: {top_quartile}"
        )

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):
    import re
    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: Assign Performance Quartiles

解答