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_columns_present(self):
        """Check that 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]
        expected_columns = ["sale_id", "sale_date", "employee_id", "amount", "moving_avg_7d"]
        missing = [col for col in expected_columns if col not in columns]
        _dynamic_test(
            self,
            len(missing) == 0,
            "All required columns are present in the output.",
            f"Missing columns: {missing}"
        )

    def test_ordering(self):
        """Check that the results are ordered by employee_id and sale_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)
        if not df.empty:
            sorted_df = df.sort_values(["employee_id", "sale_date"]).reset_index(drop=True)
            condition = df[["employee_id", "sale_date"]].reset_index(drop=True).equals(
                sorted_df[["employee_id", "sale_date"]]
            )
        else:
            condition = True
        _dynamic_test(
            self,
            condition,
            "Results are ordered by employee_id and sale_date.",
            "Results are not ordered by employee_id and sale_date."
        )

    def test_moving_average_correctness(self):
        """Check the correctness of the 7-day moving average calculation for a sample employee."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Get user result
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        df = pd.DataFrame(rows, columns=columns)
        # Check only for employee_id = 101
        df_101 = df[df["employee_id"] == 101].sort_values("sale_date").reset_index(drop=True)
        # Get correct moving averages for employee 101
        amounts = [1200.50, 1100.00, 1250.00, 1150.00, 1280.00, 1190.00]
        dates = ["2024-06-01", "2024-06-02", "2024-06-04", "2024-06-06", "2024-06-08", "2024-06-10"]
        # Manually calculate moving averages for each row for employee 101
        sale_dates = ["2024-06-01", "2024-06-02", "2024-06-04", "2024-06-06", "2024-06-08", "2024-06-10"]
        amounts = [1200.50, 1100.00, 1250.00, 1150.00, 1280.00, 1190.00]
        expected_avgs = []
        for i in range(len(amounts)):
            window = amounts[max(0, i-6):i+1]
            avg = round(sum(window)/len(window), 5)
            expected_avgs.append(avg)
        # Compare with user result (allow small rounding difference)
        correct = True
        for idx, row in enumerate(df_101.itertuples()):
            # Find corresponding date in manual calculation
            if str(row.sale_date) in sale_dates:
                i = sale_dates.index(str(row.sale_date))
                user_avg = float(row.moving_avg_7d)
                if not abs(user_avg - expected_avgs[i]) < 0.01:
                    correct = False
        _dynamic_test(
            self,
            correct,
            "Moving average is correctly calculated for employee 101.",
            "Moving average is not correctly calculated for employee 101."
        )

    def test_partition_by_employee(self):
        """Check that moving average is calculated separately for each employee."""
        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)
        # For each employee, check that moving_avg_7d is not affected by other employees
        affected = False
        for emp_id in df["employee_id"].unique():
            emp_df = df[df["employee_id"] == emp_id].sort_values("sale_date").reset_index(drop=True)
            # For each row, moving average must only use this employee's rows
            for idx, row in enumerate(emp_df.itertuples()):
                # The moving average should be between min and max of this employee's amount so far
                window = emp_df.iloc[max(0, idx-6):idx+1]["amount"].astype(float)
                avg = round(window.mean(), 5)
                if not abs(float(row.moving_avg_7d) - avg) < 0.01:
                    affected = True
        _dynamic_test(
            self,
            not affected,
            "Moving average is calculated separately for each employee.",
            "Moving average is not partitioned by employee_id."
        )

    def test_window_frame(self):
        """Check that the window frame includes the current row and previous 6 rows only."""
        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)
        # Pick employee 102 (has more than 7 sales)
        emp_id = 102
        emp_df = df[df["employee_id"] == emp_id].sort_values("sale_date").reset_index(drop=True)
        # For each row, check window size is at most 7
        correct = True
        for idx, row in enumerate(emp_df.itertuples()):
            window = emp_df.iloc[max(0, idx-6):idx+1]["amount"].astype(float)
            avg = round(window.mean(), 5)
            if not abs(float(row.moving_avg_7d) - avg) < 0.01:
                correct = False
        _dynamic_test(
            self,
            correct,
            "Window frame includes current row and previous 6 rows only.",
            "Window frame is not correct (not current row and previous 6 rows)."
        )


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: 7-Day Moving Average of Sales

解答