import unittest
import psycopg2
import pandas as pd
import re

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 setUp(self):
        try:
            self.cursor.execute("ROLLBACK;")
        except Exception:
            pass
        self.connection.commit()
        # Only add signup_date column if it does not exist
        self.cursor.execute("""
            SELECT column_name FROM information_schema.columns
            WHERE table_name = 'customers' AND column_name = 'signup_date';
        """)
        exists = self.cursor.fetchone()
        if not exists:
            self.cursor.execute("ALTER TABLE customers ADD COLUMN signup_date DATE;")
            self.connection.commit()
        # Always update signup dates regardless
        self.cursor.execute("""
            UPDATE customers SET signup_date = '2021-03-15' WHERE customer_id = 1;
            UPDATE customers SET signup_date = '2022-07-22' WHERE customer_id = 2;
            UPDATE customers SET signup_date = '2022-12-05' WHERE customer_id = 3;
            UPDATE customers SET signup_date = '2023-01-18' WHERE customer_id = 4;
            UPDATE customers SET signup_date = '2023-06-09' WHERE customer_id = 5;
            UPDATE customers SET signup_date = '2023-09-27' WHERE customer_id = 6;
            UPDATE customers SET signup_date = '2024-02-14' WHERE customer_id = 7;
            UPDATE customers SET signup_date = '2024-04-02' WHERE customer_id = 8;
        """)
        self.connection.commit()

    def test_group_by_signup_year(self):
        """Check if the query groups customers by signup year and counts signups correctly."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Remove any ALTER TABLE statement from the query
        query = re.sub(r"(?im)^.*alter table customers add column signup_date date;.*$", "", query)
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        df = pd.DataFrame(rows, columns=columns)
        expected = pd.DataFrame({
            "signup_year": [2021.0, 2022.0, 2023.0, 2024.0],
            "signup_count": [1, 2, 3, 2]
        })
        # Check that all expected years and counts are present
        merged = pd.merge(df, expected, on=["signup_year", "signup_count"], how="outer", indicator=True)
        missing = merged[merged['_merge'] != 'both']
        condition = missing.empty
        _dynamic_test(
            self,
            condition,
            "Query groups and counts signups by year correctly.",
            f"Query result missing or incorrect for signup years/counts: {missing.to_dict(orient='records')}"
        )

    def test_one_row_per_year(self):
        """Check that only one row per unique signup year is returned."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Remove any ALTER TABLE statement from the query
        query = re.sub(r"(?im)^.*alter table customers add column signup_date date;.*$", "", query)
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        df = pd.DataFrame(rows, columns=columns)
        years = df['signup_year'].tolist()
        unique_years = set(years)
        counts = df['signup_year'].value_counts()
        condition = all(counts[year] == 1 for year in unique_years)
        _dynamic_test(
            self,
            condition,
            "Only one row per unique signup year.",
            "There are duplicate rows for some signup years."
        )

    def test_columns_and_order(self):
        """Check that columns are named correctly and ordered by signup_year ascending."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Remove any ALTER TABLE statement from the query
        query = re.sub(r"(?im)^.*alter table customers add column signup_date date;.*$", "", query)
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        df = pd.DataFrame(rows, columns=columns)
        # Check column names
        condition_cols = (normalize_text(columns[0]) == "signup_year" and normalize_text(columns[1]) == "signup_count")
        # Check order ascending
        condition_order = (df['signup_year'].tolist() == sorted(df['signup_year'].tolist()))
        _dynamic_test(
            self,
            condition_cols and condition_order,
            "Columns are named correctly and results are ordered by signup_year ascending.",
            "Columns are misnamed or results are not ordered by signup_year ascending."
        )

    def test_grouping_uses_extract(self):
        """Check that EXTRACT(YEAR FROM signup_date) is used in SELECT or GROUP BY."""
        with open("query.sql", "r") as file:
            query = file.read().lower()
        uses_extract = "extract(year from signup_date)" in query
        _dynamic_test(
            self,
            uses_extract,
            "Query uses EXTRACT(YEAR FROM signup_date) to group by year.",
            "Query does not use EXTRACT(YEAR FROM signup_date) to group by year."
        )

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

A comprehensive beginner-friendly course designed to teach SQL from the ground up, focusing on practical data analysis skills. Through engaging explanations, real-world examples, and hands-on challenges, learners will master querying, transforming, and analyzing data using SQL.

Learn the foundational concepts of relational databases and SQL. Explore how to retrieve, filter, and sort data from tables, setting the stage for deeper data analysis.

Dive into SQL's powerful aggregation features. Learn to summarize, group, and analyze data for deeper insights.

Unlock the power of combining data from multiple tables. Learn to perform different types of joins to answer complex business questions.

Master advanced SQL techniques for powerful data analysis, including subqueries, window functions, and data transformation.

Challenge: Customer Signup Years

Solution

Challenge: Customer Signup Years

Solution