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_and_types(self):
        """Check if required columns and types are present."""
        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 = ['ab_group','user_id','event_count','user_rank']
        missing = [col for col in expected_columns if col not in columns]
        _dynamic_test(
            self,
            len(missing) == 0,
            "All required columns are present.",
            f"Missing columns: {', '.join(missing)}"
        )

    def test_ranking_within_ab_group(self):
        """Check that users are ranked correctly by event_count within each ab_group."""
        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 ab_group, check rank order
        passed = True
        failure_msgs = []
        for grp in df['ab_group'].unique():
            sub = df[df['ab_group'] == grp].sort_values('user_rank')
            counts = sub['event_count'].tolist()
            ranks = sub['user_rank'].tolist()
            # Check that event_count is non-increasing as rank increases
            if any(counts[i] < counts[i+1] for i in range(len(counts)-1)):
                passed = False
                failure_msgs.append(f"In ab_group {grp}, event_count is not sorted by rank.")
            # Check that rank starts at 1
            if len(ranks) > 0 and ranks[0] != 1:
                passed = False
                failure_msgs.append(f"In ab_group {grp}, rank does not start at 1.")
        _dynamic_test(
            self,
            passed,
            "Ranking and ordering within ab_group is correct.",
            " ".join(failure_msgs)
        )

    def test_event_count_aggregation(self):
        """Check that event_count matches the number of events per user."""
        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)
        # Calculate ground truth
        self.cursor.execute("SELECT user_id, COUNT(event_id) FROM events GROUP BY user_id")
        truth = dict(self.cursor.fetchall())
        df['event_count'] = df['event_count'].astype(int)
        mismatches = []
        for _, row in df.iterrows():
            expected = truth.get(row['user_id'], 0)
            if row['event_count'] != expected:
                mismatches.append((row['user_id'], row['event_count'], expected))
        _dynamic_test(
            self,
            len(mismatches) == 0,
            "event_count matches the number of events per user.",
            f"event_count mismatches: {mismatches}"
        )

    def test_all_users_included(self):
        """Check that all users are present in the output."""
        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)
        self.cursor.execute("SELECT COUNT(*) FROM users")
        user_count = self.cursor.fetchone()[0]
        _dynamic_test(
            self,
            len(df) == user_count,
            "All users are present in the output.",
            f"Output has {len(df)} rows, expected {user_count}."
        )

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 SQL from the ground up with a focus on real-world product analytics. Learn to analyze sessionized data, measure retention and churn, track feature adoption, and perform A/B test analysis using advanced SQL techniques. This course is designed for aspiring product analysts with no prior SQL experience, blending engaging theory with hands-on challenges.

Learn how to structure, query, and analyze sessionized user data, a foundational skill for product analytics.

Dive into retention and churn metrics, learning how to measure and interpret user engagement over time.

Learn to track and analyze how users discover and adopt new product features using SQL.

Master advanced SQL techniques for analyzing A/B tests and using window functions for product metrics.

Challenge: Advanced Window Analytics

解答