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_signup_month_grouping(self):
        """Check query groups by employee hire month as signup_month."""
        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)

        # All signup_month values must match employee hire months truncated to month
        self.cursor.execute("SELECT DISTINCT DATE_TRUNC('month', hire_date) FROM employees")
        expected_months = set(row[0] for row in self.cursor.fetchall())
        actual_months = set(result_df['signup_month'])
        missing = expected_months - actual_months
        _dynamic_test(
            self,
            len(missing) == 0,
            "All signup_month values are present and grouped by hire month.",
            f"Missing signup_months in result: {missing}"
        )

    def test_activity_month_identification(self):
        """Check query identifies activity_month for each performance review."""
        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 DATE_TRUNC('month', review_date) FROM performance")
        expected_months = set(row[0] for row in self.cursor.fetchall())
        actual_months = set(result_df['activity_month'])
        missing = expected_months - actual_months
        _dynamic_test(
            self,
            len(missing) == 0,
            "All activity_month values are present and identified correctly.",
            f"Missing activity_months in result: {missing}"
        )

    def test_retention_rate_calculation(self):
        """Check retention rate is correct for a known cohort and month."""
        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)

        # Instead of a hardcoded date, find a cohort with only one employee and a month where they had a review
        self.cursor.execute("""
            SELECT e.employee_id, DATE_TRUNC('month', e.hire_date) AS signup_month, DATE_TRUNC('month', p.review_date) AS activity_month
            FROM employees e
            JOIN performance p ON e.employee_id = p.employee_id
        """)
        df = pd.DataFrame(self.cursor.fetchall(), columns=["employee_id","signup_month","activity_month"])
        # Find a cohort with only one employee
        self.cursor.execute("SELECT DATE_TRUNC('month', hire_date) as signup_month, COUNT(*) FROM employees GROUP BY DATE_TRUNC('month', hire_date)")
        cohort_sizes = dict(self.cursor.fetchall())
        single_cohorts = [k for k,v in cohort_sizes.items() if v==1]
        found = False
        for signup_month in single_cohorts:
            filtered = df[df['signup_month']==signup_month]
            if not filtered.empty:
                # Take first activity_month for this cohort
                activity_month = filtered.iloc[0]['activity_month']
                user_filtered = result_df[(result_df['signup_month'] == signup_month) & (result_df['activity_month'] == activity_month)]
                if not user_filtered.empty:
                    rate = float(user_filtered.iloc[0]['retention_rate'])
                    _dynamic_test(
                        self,
                        abs(rate - 100.0) < 0.01,
                        "Retention rate for single-employee cohort and active month is 100.0.",
                        f"Expected 100.0 retention rate, got {rate}"
                    )
                    found = True
                    break
        if not found:
            _dynamic_test(
                self,
                False,
                "",
                "No single-employee cohort with a review found for testing."
            )

    def test_decimal_precision(self):
        """Check retention_rate is returned as a decimal value (not integer)."""
        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)
        # There must be at least one value with a decimal point
        has_decimal = any((isinstance(val, float) and not float(val).is_integer()) for val in result_df['retention_rate'])
        _dynamic_test(
            self,
            has_decimal or any('.' in str(val) for val in result_df['retention_rate']),
            "At least one retention_rate value has decimal precision.",
            "No retention_rate values have decimal precision."
        )

    def test_all_signup_activity_combinations_present(self):
        """Check all signup_month and activity_month combinations where cohort members were active are present."""
        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)

        # For each employee, get their signup month and each month they had a review
        self.cursor.execute("""
            SELECT DISTINCT DATE_TRUNC('month', e.hire_date) AS signup_month, DATE_TRUNC('month', p.review_date) AS activity_month
            FROM employees e
            JOIN performance p ON e.employee_id = p.employee_id
        """)
        expected = set(self.cursor.fetchall())
        actual = set(zip(result_df['signup_month'], result_df['activity_month']))
        missing = expected - actual
        _dynamic_test(
            self,
            len(missing) == 0,
            "All signup_month and activity_month combinations with active cohort members are present.",
            f"Missing combinations: {missing}"
        )


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: Monthly Retention Rate Calculation

解答