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_returns_only_customers_with_total_quantity_greater_than_5(self):
        """Check that only customers with total_quantity > 5 are returned."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        df = pd.DataFrame(rows, columns=columns)
        # All total_quantity must be > 5
        condition = True
        if not df.empty:
            condition = (df['total_quantity'] > 5).all()
        _dynamic_test(
            self,
            condition,
            "All customers in result have total_quantity greater than 5.",
            "Some customers have total_quantity less than or equal to 5."
        )

    def test_output_columns(self):
        """Check that output contains customer_id and total_quantity columns."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        columns = [desc[0] for desc in self.cursor.description]
        condition = set(columns) == set(["customer_id", "total_quantity"])
        _dynamic_test(
            self,
            condition,
            "Output contains both customer_id and total_quantity columns.",
            f"Output columns are incorrect: found {columns}."
        )

    def test_total_quantity_is_sum(self):
        """Check that total_quantity is correct sum for each customer_id."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        df = pd.DataFrame(rows, columns=columns)
        # Calculate expected sums for all customers
        self.cursor.execute("SELECT customer_id, SUM(quantity) AS total_quantity FROM sales GROUP BY customer_id HAVING SUM(quantity) > 5;")
        expected_rows = self.cursor.fetchall()
        expected_df = pd.DataFrame(expected_rows, columns=["customer_id", "total_quantity"])
        # Merge on customer_id and compare total_quantity
        merged = pd.merge(df, expected_df, on="customer_id", suffixes=("_user", "_expected"), how="outer", indicator=True)
        correct = (merged['total_quantity_user'] == merged['total_quantity_expected']).all() and (merged['_merge'] == 'both').all()
        _dynamic_test(
            self,
            correct,
            "total_quantity values are correct for all qualifying customers.",
            "total_quantity values are incorrect or missing for some customers."
        )

    def test_customers_with_quantity_5_or_less_not_included(self):
        """Check that customers with total quantity 5 or less are not included."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        df = pd.DataFrame(rows, columns=columns)
        # Find all customers with total quantity <= 5
        self.cursor.execute("SELECT customer_id, SUM(quantity) FROM sales GROUP BY customer_id HAVING SUM(quantity) <= 5;")
        low_qty_customers = [row[0] for row in self.cursor.fetchall()]
        if df.empty:
            condition = True
        else:
            condition = not df['customer_id'].isin(low_qty_customers).any()
        _dynamic_test(
            self,
            condition,
            "No customers with total quantity <= 5 are included.",
            "Some customers with total quantity <= 5 are present in the result."
        )

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: High-Volume Customers

解答