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 setUp(self):
        # Restore all tables to known state before each test
        self.cursor.execute("DELETE FROM assignments;")
        self.cursor.execute("DELETE FROM projects;")
        self.cursor.execute("DELETE FROM employees;")

        self.cursor.execute("""
            INSERT INTO employees (employee_id, name, manager_id)
            VALUES
                (1, 'Alice Johnson', NULL),
                (2, 'Bob Smith', 1),
                (3, 'Carol White', 1),
                (4, 'David Brown', 2),
                (5, 'Eve Black', 2),
                (6, 'Frank Green', 3),
                (7, 'Grace Lee', 3),
                (8, 'Hannah Adams', 1);
        """)

        self.cursor.execute("DELETE FROM departments;")
        self.cursor.execute("""
            INSERT INTO departments (department_id, department_name)
            VALUES
                (1, 'Engineering'),
                (2, 'HR'),
                (3, 'Finance'),
                (4, 'Marketing');
        """)

        self.cursor.execute("""
            INSERT INTO projects (project_id, project_name, department_id)
            VALUES
                (1, 'Website Redesign', 1),
                (2, 'Recruitment Drive', 2),
                (3, 'Budget Planning', 3),
                (4, 'Product Launch', 4);
        """)

        self.cursor.execute("""
            INSERT INTO assignments (assignment_id, employee_id, project_id)
            VALUES
                (1, 2, 1),
                (2, 4, 1),
                (3, 5, 4),
                (4, 2, 4),
                (5, 3, 2),
                (6, 6, 2),
                (7, 7, 2),
                (8, 8, 3);
        """)
        self.connection.commit()

    def test_all_assignments_present(self):
        """Check if all assignments are present with correct joins."""
        with open("query.sql", "r") as file:
            query = file.read()

        solution_query = """
            SELECT
                assignments.assignment_id,
                employees.name AS employee_name,
                projects.project_name,
                departments.department_name
            FROM assignments
            JOIN employees ON assignments.employee_id = employees.employee_id
            JOIN projects ON assignments.project_id = projects.project_id
            JOIN departments ON projects.department_id = departments.department_id
        """

        self.cursor.execute(solution_query)
        solution_rows = self.cursor.fetchall()
        solution_columns = [desc[0] for desc in self.cursor.description]
        solution_df = pd.DataFrame(solution_rows, columns=solution_columns)

        self.cursor.execute(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)

        # Normalize both dataframes safely (no applymap)
        solution_df = normalize_df(solution_df)
        user_df = normalize_df(user_df)

        # Compare by expected rows (order-independent)
        expected_cols = list(solution_df.columns)
        try:
            user_df_subset = user_df[expected_cols]
        except KeyError:
            _dynamic_test(
                self,
                False,
                "All expected assignment rows are present with correct joins.",
                "User query does not output the required columns for row comparison."
            )
            return

        expected_set = rows_as_set(solution_df, expected_cols)
        actual_set = rows_as_set(user_df_subset, expected_cols)

        missing = expected_set - actual_set
        condition = len(missing) == 0

        failure_msg = ""
        if not condition:
            failure_msg = "The following expected rows are missing:\\n" + "\\n".join(sorted(missing))

        _dynamic_test(
            self,
            condition,
            "All expected assignment rows are present with correct joins.",
            failure_msg
        )

    def test_output_columns(self):
        """Check output columns are correct and in any order."""
        with open("query.sql", "r") as file:
            query = file.read()

        self.cursor.execute(query)
        user_columns = set([normalize_text(desc[0]) for desc in self.cursor.description])

        expected_columns = set([
            "assignment_id",
            "employee_name",
            "project_name",
            "department_name"
        ])

        condition = expected_columns == user_columns
        missing = expected_columns - user_columns
        extra = user_columns - expected_columns

        failure_msg = ""
        if not condition:
            failure_msg = f"Output columns mismatch. Missing: {missing}. Extra: {extra}."

        _dynamic_test(
            self,
            condition,
            "Output includes assignment_id, employee_name, project_name, department_name.",
            failure_msg
        )

    def test_inner_joins_only(self):
        """Check only assignments with valid links in all tables are returned (inner join behavior)."""
        self.cursor.execute("DELETE FROM assignments;")
        self.connection.commit()

        try:
            self.cursor.execute("ALTER TABLE assignments DISABLE TRIGGER ALL;")
            self.cursor.execute(
                "INSERT INTO assignments (assignment_id, employee_id, project_id) VALUES (99, 100, 200);"
            )
            self.connection.commit()

            with open("query.sql", "r") as file:
                query = file.read()

            self.cursor.execute(query)
            rows = self.cursor.fetchall()

            condition = len(rows) == 0
            _dynamic_test(
                self,
                condition,
                "Query returns only assignments with valid links in all tables.",
                "Query returns assignments with missing links; should only return valid ones."
            )
        finally:
            self.cursor.execute("DELETE FROM assignments WHERE assignment_id = 99;")
            self.cursor.execute("ALTER TABLE assignments ENABLE TRIGGER ALL;")
            self.connection.commit()

        # Restore assignments for other tests
        self.cursor.execute("DELETE FROM assignments;")
        self.cursor.execute("""
            INSERT INTO assignments (assignment_id, employee_id, project_id)
            VALUES
                (1, 2, 1),
                (2, 4, 1),
                (3, 5, 4),
                (4, 2, 4),
                (5, 3, 2),
                (6, 6, 2),
                (7, 7, 2),
                (8, 8, 3);
        """)
        self.connection.commit()

    def test_joins_on_correct_columns(self):
        """Check joins use correct columns by removing a join and expecting empty output."""
        self.cursor.execute("DELETE FROM assignments;")
        self.cursor.execute("DELETE FROM employees;")
        self.connection.commit()

        with open("query.sql", "r") as file:
            query = file.read()

        self.cursor.execute(query)
        rows = self.cursor.fetchall()

        condition = len(rows) == 0
        _dynamic_test(
            self,
            condition,
            "Joins are performed on correct columns (employees missing, so output empty).",
            "Query does not join on correct columns: should return no rows if employees table is empty."
        )

        # Restore employees and assignments for other tests
        self.cursor.execute("DELETE FROM employees;")
        self.cursor.execute("""
            INSERT INTO employees (employee_id, name, manager_id)
            VALUES
                (1, 'Alice Johnson', NULL),
                (2, 'Bob Smith', 1),
                (3, 'Carol White', 1),
                (4, 'David Brown', 2),
                (5, 'Eve Black', 2),
                (6, 'Frank Green', 3),
                (7, 'Grace Lee', 3),
                (8, 'Hannah Adams', 1);
        """)
        self.cursor.execute("DELETE FROM assignments;")
        self.cursor.execute("""
            INSERT INTO assignments (assignment_id, employee_id, project_id)
            VALUES
                (1, 2, 1),
                (2, 4, 1),
                (3, 5, 4),
                (4, 2, 4),
                (5, 3, 2),
                (6, 6, 2),
                (7, 7, 2),
                (8, 8, 3);
        """)
        self.connection.commit()


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: str) -> str:
    text = str(text).lower()
    text = re.sub(r"\s{2,}", " ", text)
    text = re.sub(r"\s*([,:?])\s*", r"\1 ", text)
    return text.strip()


def normalize_df(df: pd.DataFrame) -> pd.DataFrame:
    """
    Safe replacement for DataFrame.applymap(normalize_text).
    Works even if applymap is unavailable.
    """
    # normalize column names too (helps with weird casing/spaces)
    df = df.copy()
    df.columns = [normalize_text(c) for c in df.columns]

    # cast to string then normalize each column via Series.map
    for col in df.columns:
        df[col] = df[col].astype(str).map(normalize_text)
    return df


def rows_as_set(df: pd.DataFrame, cols: list[str]) -> set[str]:
    """
    Convert rows to a set of canonical strings so comparisons are order-independent.
    """
    tmp = df[cols].copy()
    # ensure stable column order
    tmp = tmp[sorted(tmp.columns)]
    return set(tmp.apply(lambda r: "|".join(r.values.tolist()), axis=1))


if __name__ == "__main__":
    unittest.main()


test_solution.py

Master the art of SQL joins with engaging, real-world examples and hands-on challenges. This course takes you from foundational join concepts to advanced multi-table queries, ensuring you can confidently analyze and combine data from multiple sources.

Explore the essential SQL join types—INNER, LEFT, RIGHT, and FULL OUTER JOIN—through practical scenarios and hands-on challenges.

Level up your SQL skills by mastering multi-table joins, self-joins, and advanced join conditions for complex data analysis.

Master the nuances of NULL handling in joins, optimize join performance, and tackle real-world data analysis challenges.

Challenge: Complex Assignment Matching

解答