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_joins_users_and_events(self):
        """Check if the query joins users and events tables on user_id (via measurable impact on AVG)."""
        with open("query.sql", "r") as file:
            query = file.read()
    
        # Insert a new user in group A and a purchase event for them
        self.cursor.execute(
            "INSERT INTO users (signup_date, ab_group) VALUES ('2024-06-09', 'A') RETURNING user_id;"
        )
        new_user_id = self.cursor.fetchone()[0]
    
        self.cursor.execute(
            """
            INSERT INTO events (user_id, event_type, event_date, value)
            VALUES (%s, 'purchase', %s, %s);
            """,
            (new_user_id, '2024-06-09 09:00:00', 99.99)
        )
        self.connection.commit()
    
        # Run student's query
        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 the expected AVG directly from the database after the insertion
        self.cursor.execute("""
            SELECT AVG(e.value)
            FROM users u
            JOIN events e ON u.user_id = e.user_id
            WHERE u.ab_group = 'A' AND e.event_type = 'purchase';
        """)
        expected_avg = float(self.cursor.fetchone()[0])
    
        # Validate that student's result contains group A with matching AVG
        found = False
        if 'ab_group' in result_df.columns and 'avg_purchase_value' in result_df.columns:
            a_row = result_df[result_df['ab_group'] == 'A']
            if not a_row.empty:
                actual_avg = float(a_row.iloc[0]['avg_purchase_value'])
                found = abs(actual_avg - expected_avg) < 0.01
    
        _dynamic_test(
            self,
            found,
            "Query joins users and events tables on user_id.",
            "Query does not properly join users and events tables on user_id."
        )



    def test_filters_purchase_events(self):
        """Check if the query filters only 'purchase' events."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Insert a non-purchase event with a high value for a user in group B
        self.cursor.execute("INSERT INTO users (signup_date, ab_group) VALUES ('2024-06-10', 'B') RETURNING user_id;")
        new_user_id = self.cursor.fetchone()[0]
        self.cursor.execute(f"INSERT INTO events (user_id, event_type, event_date, value) VALUES ({new_user_id}, 'refund', '2024-06-10 10:00:00', 999.99);")
        self.connection.commit()
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(rows, columns=columns)
        # The high value from 'refund' should not affect avg_purchase_value
        # Compare to control: avg_purchase_value for group B should not be 999.99
        if 'ab_group' in result_df.columns and 'avg_purchase_value' in result_df.columns:
            b_rows = result_df[result_df['ab_group']=='B']
            if not b_rows.empty:
                avg_val = float(b_rows.iloc[0]['avg_purchase_value'])
                condition = avg_val != 999.99
            else:
                condition = True
        else:
            condition = False
        _dynamic_test(
            self,
            condition,
            "Query filters only 'purchase' events.",
            "Query does not filter only 'purchase' events."
        )

    def test_groups_by_ab_group(self):
        """Check if the query groups results by 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]
        result_df = pd.DataFrame(rows, columns=columns)
        # There should be one row per ab_group present in the data
        ab_groups = set(result_df['ab_group']) if 'ab_group' in result_df.columns else set()
        self.cursor.execute("SELECT DISTINCT ab_group FROM users")
        expected_groups = set([r[0] for r in self.cursor.fetchall()])
        expected_groups = set(['A', 'B']) & expected_groups
        condition = ab_groups == expected_groups
        _dynamic_test(
            self,
            condition,
            "Query groups results by ab_group.",
            "Query does not group results by ab_group."
        )

    def test_calculates_average_value(self):
        """Check if the query calculates the average value for 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]
        result_df = pd.DataFrame(rows, columns=columns)
        # Calculate expected averages manually for each group
        self.cursor.execute("""
            SELECT u.ab_group, AVG(e.value) AS avg_purchase_value
            FROM users u JOIN events e ON u.user_id = e.user_id
            WHERE e.event_type = 'purchase'
            GROUP BY u.ab_group
        """)
        expected = pd.DataFrame(self.cursor.fetchall(), columns=['ab_group', 'avg_purchase_value'])
        merged = pd.merge(result_df, expected, on='ab_group', suffixes=('_user', '_expected'))
        condition = True
        for _, row in merged.iterrows():
            # Allow small floating point differences
            if abs(float(row['avg_purchase_value_user']) - float(row['avg_purchase_value_expected'])) > 0.01:
                condition = False
                break
        _dynamic_test(
            self,
            condition,
            "Query calculates average value for each ab_group.",
            "Query does not correctly calculate average value for each ab_group."
        )

    def test_orders_by_average_desc(self):
        """Check if the query orders results by average value in descending order."""
        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)
        if 'avg_purchase_value' in result_df.columns:
            avg_vals = list(result_df['avg_purchase_value'])
            sorted_vals = sorted(avg_vals, reverse=True)
            condition = avg_vals == sorted_vals
        else:
            condition = False
        _dynamic_test(
            self,
            condition,
            "Query orders results by average value in descending order.",
            "Query does not order results by average value in descending order."
        )

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: A/B Test Metrics

Solution

Challenge: A/B Test Metrics

Solution