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_total_sales_per_product(self):
        """Check that total sales per product is correct."""
        with open("query.sql", "r") as file:
            user_code = file.read()
        queries = [q.strip() for q in user_code.split(';') if q.strip()]
        # First query: total sales per product
        if len(queries) < 1:
            _dynamic_test(self, False, "First query present.", "First query (total sales per product) is missing.")
            return
        user_query = queries[0]
        self.cursor.execute(user_query)
        user_rows = self.cursor.fetchall()
        user_columns = [desc[0] for desc in self.cursor.description]
        user_df = pd.DataFrame(user_rows, columns=user_columns)
        # Solution query
        self.cursor.execute("""
            SELECT product_id, SUM(total_amount) AS total_sales
            FROM sales
            GROUP BY product_id
        """)
        sol_rows = self.cursor.fetchall()
        sol_columns = [desc[0] for desc in self.cursor.description]
        sol_df = pd.DataFrame(sol_rows, columns=sol_columns)
        # Check all product_ids are present and amounts match
        merged = pd.merge(sol_df, user_df, on="product_id", suffixes=("_sol","_user"), how="left")
        missing = merged[merged["total_sales_user"].isnull()]
        wrong = merged[(merged["total_sales_user"].notnull()) & (abs(merged["total_sales_sol"] - merged["total_sales_user"]) > 0.01)]
        condition = (missing.empty and wrong.empty and len(user_df) == len(sol_df))
        fail_msg = ""
        if not missing.empty:
            fail_msg += f"Missing product_ids: {missing['product_id'].tolist()}\\n"
        if not wrong.empty:
            fail_msg += f"Incorrect total_sales for product_ids: {wrong['product_id'].tolist()}\\n"
        if len(user_df) != len(sol_df):
            fail_msg += f"Incorrect number of rows. Expected {len(sol_df)}, got {len(user_df)}.\\n"
        _dynamic_test(
            self,
            condition,
            "Total sales per product is correct for all product_ids.",
            fail_msg.strip()
        )

    def test_order_count_per_customer(self):
        """Check that order count per customer is correct."""
        with open("query.sql", "r") as file:
            user_code = file.read()
        queries = [q.strip() for q in user_code.split(';') if q.strip()]
        # Second query: order count per customer
        if len(queries) < 2:
            _dynamic_test(self, False, "Second query present.", "Second query (order count per customer) is missing.")
            return
        user_query = queries[1]
        self.cursor.execute(user_query)
        user_rows = self.cursor.fetchall()
        user_columns = [desc[0] for desc in self.cursor.description]
        user_df = pd.DataFrame(user_rows, columns=user_columns)
        # Solution query
        self.cursor.execute("""
            SELECT customer_id, COUNT(sale_id) AS order_count
            FROM sales
            GROUP BY customer_id
        """)
        sol_rows = self.cursor.fetchall()
        sol_columns = [desc[0] for desc in self.cursor.description]
        sol_df = pd.DataFrame(sol_rows, columns=sol_columns)
        merged = pd.merge(sol_df, user_df, on="customer_id", suffixes=("_sol","_user"), how="left")
        missing = merged[merged["order_count_user"].isnull()]
        wrong = merged[(merged["order_count_user"].notnull()) & (merged["order_count_sol"] != merged["order_count_user"])]
        condition = (missing.empty and wrong.empty and len(user_df) == len(sol_df))
        fail_msg = ""
        if not missing.empty:
            fail_msg += f"Missing customer_ids: {missing['customer_id'].tolist()}\\n"
        if not wrong.empty:
            fail_msg += f"Incorrect order_count for customer_ids: {wrong['customer_id'].tolist()}\\n"
        if len(user_df) != len(sol_df):
            fail_msg += f"Incorrect number of rows. Expected {len(sol_df)}, got {len(user_df)}.\\n"
        _dynamic_test(
            self,
            condition,
            "Order count per customer is correct for all customer_ids.",
            fail_msg.strip()
        )

    def test_group_by_used(self):
        """Check both queries use GROUP BY."""
        with open("query.sql", "r") as file:
            user_code = file.read().lower()
        queries = [q.strip() for q in user_code.split(';') if q.strip()]
        has_group_by_first = 'group by' in queries[0] if len(queries) > 0 else False
        has_group_by_second = 'group by' in queries[1] if len(queries) > 1 else False
        _dynamic_test(
            self,
            has_group_by_first and has_group_by_second,
            "Both queries use GROUP BY.",
            "Both queries must use GROUP BY."
        )

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 business intelligence. Learn to write powerful queries for KPIs, implement business logic, analyze trends with window functions, design data models for BI tools, and apply reporting patterns. Each section blends engaging theory with hands-on, real-world challenges to prepare you for BI analyst roles.

Learn to use SQL aggregations to calculate and interpret key business metrics, such as revenue, customer counts, and averages, essential for BI analysis.

Implement business rules and conditional logic in SQL queries using CASE WHEN, enabling dynamic categorization and decision-making.

Analyze business trends and rankings over time using advanced SQL window functions, essential for BI trend analysis and leaderboards.

Challenge: Segment Sales Data

Solution