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 if required columns are present in the result."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        columns = [desc[0] for desc in self.cursor.description]
        required_cols = {"cohort_start", "week_number", "retained_users"}
        condition = required_cols.issubset(set(columns))
        _dynamic_test(
            self,
            condition,
            "All required columns are present in the result.",
            f"Missing columns: {required_cols - set(columns)}"
        )

    def test_result_values(self):
        """Check if the result matches expected cohort/week/user counts."""
        with open("query.sql", "r") as file:
            user_query = file.read()
        # Reference solution
        solution_query = '''
        SELECT
          u.signup_date AS cohort_start,
          FLOOR(EXTRACT(EPOCH FROM (l.login_date::timestamp - u.signup_date::timestamp)) / (7*24*60*60)) AS week_number,
          COUNT(DISTINCT l.user_id) AS retained_users
        FROM users u
        JOIN logins l
          ON u.user_id = l.user_id
          AND l.login_date >= u.signup_date
        GROUP BY u.signup_date, week_number
        ORDER BY cohort_start, week_number;
        '''
        self.cursor.execute(solution_query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        solution_df = pd.DataFrame(rows, columns=columns)

        self.cursor.execute(user_query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        user_df = pd.DataFrame(rows, columns=columns)

        # Compare expected rows
        missing_rows = []
        for _, expected_row in solution_df.iterrows():
            expected_dict = expected_row.to_dict()
            mask = (user_df[list(expected_dict.keys())] == pd.Series(expected_dict)).all(axis=1)
            if not mask.any():
                missing_rows.append(expected_dict)
        condition = len(missing_rows) == 0
        if not condition:
            failure_msg = f"The following expected rows are missing:\\n" + "\\n".join(str(r) for r in missing_rows)
        else:
            failure_msg = ""
        _dynamic_test(
            self,
            condition,
            "All expected cohort/week rows are present with correct retained_users count.",
            failure_msg
        )

    def test_no_extra_rows(self):
        """Check that user result does not contain extra rows not in the reference solution."""
        with open("query.sql", "r") as file:
            user_query = file.read()
        # Reference solution
        solution_query = '''
        SELECT
          u.signup_date AS cohort_start,
          FLOOR(EXTRACT(EPOCH FROM (l.login_date::timestamp - u.signup_date::timestamp)) / (7*24*60*60)) AS week_number,
          COUNT(DISTINCT l.user_id) AS retained_users
        FROM users u
        JOIN logins l
          ON u.user_id = l.user_id
          AND l.login_date >= u.signup_date
        GROUP BY u.signup_date, week_number
        ORDER BY cohort_start, week_number;
        '''
        self.cursor.execute(solution_query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        solution_df = pd.DataFrame(rows, columns=columns)

        self.cursor.execute(user_query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        user_df = pd.DataFrame(rows, columns=columns)

        # Compare extra rows
        extra_rows = []
        for _, user_row in user_df.iterrows():
            user_dict = user_row.to_dict()
            mask = (solution_df[list(user_dict.keys())] == pd.Series(user_dict)).all(axis=1)
            if not mask.any():
                extra_rows.append(user_dict)
        condition = len(extra_rows) == 0
        if not condition:
            failure_msg = f"The following extra rows are present in your result:\\n" + "\\n".join(str(r) for r in extra_rows)
        else:
            failure_msg = ""
        _dynamic_test(
            self,
            condition,
            "No extra rows are present in the result.",
            failure_msg
        )

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 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: Retention Matrix

Solution

Challenge: Retention Matrix

Solution