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 test_grouping_by_year_and_month(self):
        """Check if query groups by both year and month from order_date."""
        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]
        result_df = pd.DataFrame(rows, columns=columns)

        # Should have columns for year and month
        year_col = [col for col in result_df.columns if normalize_text(col) in ["order_year", "year"]]
        month_col = [col for col in result_df.columns if normalize_text(col) in ["order_month", "month"]]
        condition = len(year_col) == 1 and len(month_col) == 1
        _dynamic_test(
            self,
            condition,
            "Query groups by both year and month.",
            "Query does not group by both year and month or columns are missing."
        )

    def test_correct_order_counts(self):
        """Check if query returns correct count of orders for each year and month."""
        with open("query.sql", "r") as file:
            query = file.read()
        solution_query = '''
            SELECT EXTRACT(YEAR FROM order_date) AS order_year,
                   EXTRACT(MONTH FROM order_date) AS order_month,
                   COUNT(*) AS order_count
            FROM orders
            GROUP BY EXTRACT(YEAR FROM order_date), EXTRACT(MONTH FROM order_date)
            ORDER BY order_year, order_month
        '''
        self.cursor.execute(solution_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)

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

        # Normalize columns for comparison
        def get_col(df, names):
            for n in names:
                for col in df.columns:
                    if normalize_text(col) == n:
                        return col
            return None
        year_col = get_col(result_df, ["order_year", "year"])
        month_col = get_col(result_df, ["order_month", "month"])
        count_col = get_col(result_df, ["order_count", "count"])
        if not year_col or not month_col or not count_col:
            _dynamic_test(
                self,
                False,
                "",
                "Query result missing required columns."
            )
            return
        # Compare rows
        missing = []
        for _, row in expected_df.iterrows():
            mask = (
                (result_df[year_col] == row["order_year"]) &
                (result_df[month_col] == row["order_month"]) &
                (result_df[count_col] == row["order_count"])
            )
            if not mask.any():
                missing.append({"year": row["order_year"], "month": row["order_month"], "order_count": row["order_count"]})
        condition = len(missing) == 0
        fail_msg = "Missing or incorrect order counts for: " + str(missing) if missing else ""
        _dynamic_test(
            self,
            condition,
            "Query returns correct order counts for each year and month.",
            fail_msg
        )

    def test_all_months_present(self):
        """Check if query includes all months present in the data."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Expected months: 2024-05 and 2024-06
        expected = set([(2024.0, 5.0), (2024.0, 6.0)])
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(rows, columns=columns)
        year_col = [col for col in result_df.columns if normalize_text(col) in ["order_year", "year"]][0]
        month_col = [col for col in result_df.columns if normalize_text(col) in ["order_month", "month"]][0]
        found = set(zip(result_df[year_col], result_df[month_col]))
        missing = expected - found
        condition = len(missing) == 0
        _dynamic_test(
            self,
            condition,
            "Query includes all months present in the data.",
            f"Missing months in result: {missing}"
        )

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

This beginner course focuses on the essential SQL functions used for data cleaning and transformation. Learners will practice handling missing values, formatting and standardizing data, and applying string, date, numeric, logical, and null handling functions to prepare datasets for analysis. The course provides hands-on challenges directly related to the most common data preparation tasks in SQL.

Challenge: Monthly Order Summary

解答