import unittest
import psycopg2
import pandas as pd
from datetime import date, timedelta

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()
        # Create performance table and sample data
        cls.cursor.execute("DROP TABLE IF EXISTS performance;")
        cls.cursor.execute("DROP TABLE IF EXISTS employees;")
        cls.cursor.execute("""
            CREATE TABLE employees (
                employee_id SERIAL PRIMARY KEY,
                name VARCHAR(100) NOT NULL,
                department VARCHAR(50) NOT NULL,
                hire_date DATE NOT NULL,
                salary DECIMAL(10,2) NOT NULL
            );
        """)
        cls.cursor.execute("""
            CREATE TABLE performance (
                performance_id SERIAL PRIMARY KEY,
                employee_id INTEGER NOT NULL REFERENCES employees(employee_id),
                review_date DATE NOT NULL,
                score INTEGER NOT NULL CHECK (score BETWEEN 1 AND 10)
            );
        """)
        cls.cursor.execute("""
            INSERT INTO employees (name, department, hire_date, salary) VALUES
            ('Alice Johnson', 'Engineering', '2021-03-15', 85000.00),
            ('Bob Smith', 'Marketing', '2020-07-22', 65000.00),
            ('Carol Lee', 'Human Resources', '2019-11-01', 60000.00);
        """)
        # Insert test data for three employees
        cls.cursor.execute("""
            INSERT INTO performance (employee_id, review_date, score) VALUES
            (1, '2022-03-20', 8),
            (1, '2023-03-18', 9),
            (2, '2021-08-01', 7),
            (2, '2022-08-01', 8),
            (2, '2023-08-01', 7),
            (3, '2020-12-15', 6),
            (3, '2021-12-15', 7);
        """)
        cls.connection.commit()

    @classmethod
    def tearDownClass(cls):
        cls.cursor.execute("DROP TABLE IF EXISTS performance;")
        cls.cursor.execute("DROP TABLE IF EXISTS employees;")
        cls.connection.commit()
        cls.cursor.close()
        cls.connection.close()

    def test_returns_one_row_per_review(self):
        """Check the query returns one row per review."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute("SELECT COUNT(*) FROM performance;")
        expected_count = self.cursor.fetchone()[0]
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        _dynamic_test(
            self,
            len(rows) == expected_count,
            "The query returns one row for each review.",
            f"Expected {expected_count} rows, got {len(rows)}."
        )

    def test_output_columns(self):
        """Check the output includes the required columns."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        columns = [desc[0] for desc in self.cursor.description]
        required = ['employee_id', 'review_date', 'score', 'rolling_3_month_avg_score']
        missing = [col for col in required if col not in columns]
        _dynamic_test(
            self,
            not missing,
            "All required columns are present in the result.",
            f"Missing columns: {', '.join(missing)}."
        )

    def test_ordering(self):
        """Check the results are ordered by employee_id, review_date."""
        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)
        sorted_df = df.sort_values(["employee_id", "review_date"]).reset_index(drop=True)
        condition = df.reset_index(drop=True).equals(sorted_df)
        _dynamic_test(
            self,
            condition,
            "Result is ordered by employee_id and review_date.",
            "Result is not ordered by employee_id and review_date."
        )

    def test_rolling_3_month_avg_score(self):
        """Check the rolling 3-month average score values are correct for each row."""
        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)
        # Ensure review_date is pandas.Timestamp
        df['review_date'] = pd.to_datetime(df['review_date'])
        # Calculate expected rolling averages manually
        expected = []
        for employee_id in df['employee_id'].unique():
            edf = df[df['employee_id'] == employee_id].copy()
            for idx, row in edf.iterrows():
                end_date = row['review_date']
                start_date = (end_date - pd.DateOffset(months=2)).replace(day=1)
                mask = (edf['review_date'] >= start_date) & (edf['review_date'] <= end_date)
                window_avg = edf.loc[mask, 'score'].astype(float).mean()
                expected.append(window_avg)
        actual = df['rolling_3_month_avg_score'].astype(float).tolist()
        # Compare floats with tolerance
        comparison = all(abs(a - b) < 0.01 for a, b in zip(expected, actual))
        _dynamic_test(
            self,
            comparison,
            "The rolling_3_month_avg_score column contains correct averages.",
            "The rolling_3_month_avg_score column contains incorrect values."
        )

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: Rolling 3-Month Spend per Customer

解答