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_top_3_customers_by_total_quantity(self):
        """Check if the query returns the correct top 3 customers by total quantity purchased."""
        with open("query.sql", "r") as file:
            query = file.read()

        # Reference query for correct top 3 customers
        solution_query = '''
            SELECT c.customer_id, c.first_name, c.last_name, SUM(s.quantity) AS total_quantity, RANK() OVER (ORDER BY SUM(s.quantity) DESC) AS customer_rank
            FROM customers c
            JOIN sales s ON c.customer_id = s.customer_id
            GROUP BY c.customer_id, c.first_name, c.last_name
            ORDER BY customer_rank
            LIMIT 3;
        '''
        self.cursor.execute(solution_query)
        ref_rows = self.cursor.fetchall()
        ref_cols = [desc[0] for desc in self.cursor.description]
        ref_df = pd.DataFrame(ref_rows, columns=ref_cols)

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

        # Compare by customer_id and total_quantity
        missing_rows = []
        for _, ref_row in ref_df.iterrows():
            mask = (
                (df['customer_id'] == ref_row['customer_id']) &
                (df['first_name'] == ref_row['first_name']) &
                (df['last_name'] == ref_row['last_name']) &
                (df['total_quantity'] == ref_row['total_quantity']) &
                (df['customer_rank'] == ref_row['customer_rank'])
            )
            if not mask.any():
                missing_rows.append(ref_row.to_dict())
        condition = len(missing_rows) == 0 and len(df) == 3
        if not condition:
            failure_msg = f"Top 3 customers by total quantity are incorrect or missing. Missing/incorrect rows: {missing_rows}"
        else:
            failure_msg = ""
        _dynamic_test(self, condition, "Query returns correct top 3 customers by total quantity.", failure_msg)

    def test_includes_required_columns(self):
        """Check that the output includes customer_id, first_name, last_name, total_quantity, and customer_rank."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        columns = [desc[0] for desc in self.cursor.description]
        required = {'customer_id', 'first_name', 'last_name', 'total_quantity', 'customer_rank'}
        missing = required - set(columns)
        condition = len(missing) == 0
        if not condition:
            failure_msg = f"Missing required columns: {missing}"
        else:
            failure_msg = ""
        _dynamic_test(self, condition, "All required columns are present in the output.", failure_msg)

    def test_ranking_uses_window_function(self):
        """Check that a window function is used to rank customers by total quantity."""
        with open("query.sql", "r") as file:
            query = file.read().lower()
        # Look for RANK() or DENSE_RANK() or ROW_NUMBER() with OVER
        uses_window = (
            ("rank()" in query or "dense_rank()" in query or "row_number()" in query)
            and "over" in query
        )
        _dynamic_test(
            self,
            uses_window,
            "A window function is used to rank customers.",
            "No window function found for ranking customers."
        )

    def test_ranking_is_descending(self):
        """Check that ranking is applied in descending order of total quantity."""
        with open("query.sql", "r") as file:
            query = file.read().lower().replace("\\n", " ")
        # Look for ORDER BY SUM(s.quantity) DESC or similar inside OVER clause
        has_desc = "order by sum(s.quantity) desc" in query or "order by sum(quantity) desc" in query
        _dynamic_test(
            self,
            has_desc,
            "Ranking is applied in descending order of total quantity.",
            "Ranking is not applied in descending order of total quantity."
        )

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: Top 3 Customers

解答