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_row_count(self):
        """Check correct columns and only users with purchases are included."""
        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)
        expected_columns = {'user_id', 'ab_group', 'purchase_count', 'group_purchase_rank'}
        _dynamic_test(
            self,
            set(result_df.columns) == expected_columns,
            "Returned columns are correct.",
            f"Returned columns are {result_df.columns.tolist()}, but expected {list(expected_columns)}."
        )
        # Only users with at least one purchase
        self.cursor.execute("SELECT DISTINCT user_id FROM events WHERE event_type = 'purchase'")
        purchase_users = set(r[0] for r in self.cursor.fetchall())
        result_users = set(result_df['user_id'])
        _dynamic_test(
            self,
            result_users == purchase_users,
            "Only users with purchases are included.",
            f"Result includes users without purchases or misses users with purchases. Expected users: {purchase_users}, got: {result_users}."
        )

    def test_purchase_count_accuracy(self):
        """Check that purchase_count per user is correct."""
        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)
        # Compute expected purchase counts
        self.cursor.execute("""
            SELECT user_id, COUNT(*) as cnt FROM events WHERE event_type='purchase' GROUP BY user_id
        """)
        expected_counts = dict(self.cursor.fetchall())
        mismatches = []
        for _, row in result_df.iterrows():
            uid = row['user_id']
            expected = expected_counts.get(uid, 0)
            if int(row['purchase_count']) != expected:
                mismatches.append((uid, row['purchase_count'], expected))
        condition = len(mismatches) == 0
        failure_msg = "Purchase counts do not match for: " + str(mismatches) if not condition else ""
        _dynamic_test(
            self,
            condition,
            "Purchase counts per user are correct.",
            failure_msg
        )

    def test_dense_rank_within_group(self):
        """Check DENSE_RANK() within each ab_group based on purchase_count desc."""
        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)
        # Check rank logic for each group
        for group in result_df['ab_group'].unique():
            group_df = result_df[result_df['ab_group'] == group].copy()
            group_df['purchase_count'] = group_df['purchase_count'].astype(int)
            group_df = group_df.sort_values(['purchase_count', 'user_id'], ascending=[False, True])
            group_df['expected_rank'] = group_df['purchase_count'].rank(method='dense', ascending=False).astype(int)
            mismatches = group_df[group_df['expected_rank'] != group_df['group_purchase_rank']]
            condition = mismatches.empty
            failure_msg = f"Incorrect dense rank for ab_group {group}:\\n{mismatches[['user_id','purchase_count','group_purchase_rank','expected_rank']].to_dict('records')}" if not condition else ""
            _dynamic_test(
                self,
                condition,
                f"DENSE_RANK() within ab_group {group} is correct.",
                failure_msg
            )

    def test_ordering(self):
        """Check results are ordered by ab_group and group_purchase_rank."""
        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)
        ab_groups = result_df['ab_group'].tolist()
        ranks = result_df['group_purchase_rank'].tolist()
        correct_order = True
        for i in range(1, len(result_df)):
            if ab_groups[i] < ab_groups[i-1]:
                correct_order = False
                break
            if ab_groups[i] == ab_groups[i-1] and ranks[i] < ranks[i-1]:
                correct_order = False
                break
        _dynamic_test(
            self,
            correct_order,
            "Results are ordered by ab_group and group_purchase_rank.",
            "Results are not ordered by ab_group and group_purchase_rank."
        )

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: Window Function Metrics

Solution

Challenge: Window Function Metrics

Solution