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_above_average_customers_rows(self):
        """Check that only above-average customers are returned with expected fields."""
        with open("query.sql", "r") as file:
            query = file.read()

        # Calculate expected average total quantity per customer
        self.cursor.execute("""
            SELECT customer_id, SUM(quantity) AS total_quantity
            FROM sales
            GROUP BY customer_id
        """)
        rows = self.cursor.fetchall()
        customer_totals = [row[1] for row in rows]
        avg_total = sum(customer_totals) / len(customer_totals)
        above_avg_customers = [row[0] for row in rows if row[1] > avg_total]

        # Get expected result
        if above_avg_customers:
            format_ids = ','.join(str(i) for i in above_avg_customers)
            expected_query = f"""
                SELECT c.customer_id, c.first_name, c.last_name, SUM(s.quantity) AS total_quantity
                FROM customers c
                JOIN sales s ON c.customer_id = s.customer_id
                WHERE c.customer_id IN ({format_ids})
                GROUP BY c.customer_id, c.first_name, c.last_name
            """
            self.cursor.execute(expected_query)
            expected_rows = self.cursor.fetchall()
            expected_columns = [desc[0] for desc in self.cursor.description]
            expected_df = pd.DataFrame(expected_rows, columns=expected_columns)
        else:
            expected_df = pd.DataFrame(columns=["customer_id", "first_name", "last_name", "total_quantity"])

        # Run user query
        self.cursor.execute(query)
        result_rows = self.cursor.fetchall()
        result_columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(result_rows, columns=result_columns)

        # Check columns
        required_cols = ["customer_id", "first_name", "last_name", "total_quantity"]
        cond_cols = set(required_cols).issubset(set(result_df.columns))
        _dynamic_test(self, cond_cols,
            "All required columns are present in the result.",
            f"Missing one or more required columns: {set(required_cols) - set(result_df.columns)}")

        # Check expected rows are present
        missing_rows = []
        for _, exp_row in expected_df.iterrows():
            mask = (result_df[required_cols] == exp_row[required_cols]).all(axis=1)
            if not mask.any():
                missing_rows.append(exp_row[required_cols].to_dict())
        cond_rows = len(missing_rows) == 0
        fail_msg = f"The following expected rows are missing: {missing_rows}" if not cond_rows else ""
        _dynamic_test(self, cond_rows,
            "All above-average customers are present in result.",
            fail_msg)

        # Check that no below/equal average customers are present
        forbidden_ids = [row[0] for row in rows if row[1] <= avg_total]
        forbidden_present = result_df['customer_id'].isin(forbidden_ids).any()
        _dynamic_test(self, not forbidden_present,
            "No customers with total quantity less than or equal to the average are present.",
            "Result contains customers with total quantity less than or equal to the average.")

    def test_query_uses_subquery(self):
        """Check that the query uses a subquery to calculate the average."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Simple check for subquery usage
        found_subquery = "select" in query.lower() and "avg" in query.lower()
        _dynamic_test(self, found_subquery,
            "Query uses a subquery to calculate the average.",
            "Query does not use a subquery to calculate the average.")

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: Above-Average Customers

Solution