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_one_row_per_employee(self):
        """Query returns one row per employee, including those with zero reviews."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute("SELECT COUNT(*) FROM employees;")
        employee_count = self.cursor.fetchone()[0]
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        _dynamic_test(
            self,
            len(rows) == employee_count,
            "Query returns one row per employee.",
            f"Query returned {len(rows)} rows, but there are {employee_count} employees."
        )

    def test_total_score_and_review_count(self):
        """Check that total_score and review_count columns are correct for each employee."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Calculate expected values
        self.cursor.execute("""
            SELECT e.employee_id, COALESCE(SUM(p.score),0) AS total_score, COUNT(p.performance_id) AS review_count
            FROM employees e
            LEFT JOIN performance p ON e.employee_id = p.employee_id
            GROUP BY e.employee_id
            ORDER BY e.employee_id;
        """)
        expected = self.cursor.fetchall()
        self.cursor.execute(query)
        result = self.cursor.fetchall()
        # Compare total_score and review_count
        errors = []
        for i, row in enumerate(result):
            emp_id = row[0]
            expected_row = next((e for e in expected if e[0] == emp_id), None)
            if not expected_row:
                errors.append(f"Employee {emp_id} missing in result.")
                continue
            if row[2] != expected_row[1] or row[3] != expected_row[2]:
                errors.append(f"Employee {emp_id}: expected (score: {expected_row[1]}, count: {expected_row[2]}) got (score: {row[2]}, count: {row[3]})")
        _dynamic_test(
            self,
            not errors,
            "Total score and review count are correct for all employees.",
            "Errors in total_score/review_count: " + ", ".join(errors)
        )

    def test_score_quartile_and_frequency_quartile(self):
        """Check that score_quartile and frequency_quartile columns are present and correct range (1-4)."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        errors = []
        for row in rows:
            if len(row) < 6:
                errors.append(f"Row for employee_id {row[0]} does not have enough columns.")
                continue
            score_q = row[4]
            freq_q = row[5]
            if not (1 <= score_q <= 4):
                errors.append(f"Employee {row[0]}: score_quartile is {score_q}, should be 1-4.")
            if not (1 <= freq_q <= 4):
                errors.append(f"Employee {row[0]}: frequency_quartile is {freq_q}, should be 1-4.")
        _dynamic_test(
            self,
            not errors,
            "All employees have valid quartile values (1-4).",
            "Invalid quartile values: " + ", ".join(errors)
        )

    def test_zero_review_employee_included(self):
        """Employees with no reviews have total_score 0 and are included in quartiles."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Find employees with no reviews
        self.cursor.execute("""
            SELECT e.employee_id
            FROM employees e
            LEFT JOIN performance p ON e.employee_id = p.employee_id
            GROUP BY e.employee_id
            HAVING COUNT(p.performance_id) = 0;
        """)
        no_review_ids = [row[0] for row in self.cursor.fetchall()]
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        errors = []
        for emp_id in no_review_ids:
            matching = [row for row in rows if row[0] == emp_id]
            if not matching:
                errors.append(f"Employee {emp_id} with zero reviews missing in result.")
            else:
                row = matching[0]
                if row[2] != 0 or row[3] != 0:
                    errors.append(f"Employee {emp_id} with zero reviews has total_score {row[2]}, review_count {row[3]}")
                if row[4] is None or row[5] is None:
                    errors.append(f"Employee {emp_id} with zero reviews missing quartile values.")
        _dynamic_test(
            self,
            not errors,
            "Employees with zero reviews are included and have correct values.",
            "Zero-review employee errors: " + ", ".join(errors)
        )

    def test_quartile_distribution(self):
        """Employees with same score or review count are distributed into quartiles according to NTILE rules."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        # Check that quartile values are distributed as evenly as possible
        score_q_counts = {}
        freq_q_counts = {}
        for row in rows:
            score_q = row[4]
            freq_q = row[5]
            score_q_counts[score_q] = score_q_counts.get(score_q, 0) + 1
            freq_q_counts[freq_q] = freq_q_counts.get(freq_q, 0) + 1
        # For 6 employees and 4 quartiles, expected distribution: 2,1,1,2 or similar
        score_counts_sorted = sorted(score_q_counts.values(), reverse=True)
        freq_counts_sorted = sorted(freq_q_counts.values(), reverse=True)
        # Acceptable if max-min <= 1
        score_ok = max(score_q_counts.values()) - min(score_q_counts.values()) <= 1
        freq_ok = max(freq_q_counts.values()) - min(freq_q_counts.values()) <= 1
        _dynamic_test(
            self,
            score_ok and freq_ok,
            "Quartile distribution is as even as possible for both score and frequency.",
            f"Score quartile counts: {score_q_counts}, Frequency quartile counts: {freq_q_counts}"
        )


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: Customer Segmentation by Spend and Frequency

解答