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_only_null_manager_id(self):
        """Check that only employees with null manager_id 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)
        condition = df['manager_id'].isnull().all() and len(df) > 0
        _dynamic_test(
            self,
            condition,
            "All returned employees have null manager_id.",
            "Some returned employees have a non-null manager_id, or result is empty."
        )

    def test_all_columns_included(self):
        """Check that all columns from employees table are included."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        columns = [desc[0] for desc in self.cursor.description]
        expected_columns = ["employee_id", "name", "department", "manager_id", "bonus"]
        condition = set(expected_columns) == set(columns)
        _dynamic_test(
            self,
            condition,
            "All columns from employees table are included.",
            f"Missing or extra columns in result: {set(columns) ^ set(expected_columns)}"
        )

    def test_returns_correct_employees(self):
        """Check that correct employees are returned when manager_id is null."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Get expected employees
        self.cursor.execute("SELECT employee_id FROM employees WHERE manager_id IS NULL")
        expected_ids = set(row[0] for row in self.cursor.fetchall())
        # Run user query
        self.cursor.execute(query)
        result_ids = set(row[0] for row in self.cursor.fetchall())
        condition = result_ids == expected_ids
        _dynamic_test(
            self,
            condition,
            "Returned employees match those with null manager_id.",
            f"Returned employee IDs do not match: {expected_ids ^ result_ids}"
        )

    def test_works_with_additional_nulls(self):
        """Check that query works if more employees with null manager_id are added."""
        # Add a test employee with null manager_id
        self.cursor.execute("""
            INSERT INTO employees (employee_id, name, department, manager_id, bonus)
            VALUES (999, 'Test User', 'Test', NULL, 999.99)
        """)
        self.connection.commit()
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        employee_ids = [row[0] for row in rows]
        condition = 999 in employee_ids
        _dynamic_test(
            self,
            condition,
            "Query includes newly added employee with null manager_id.",
            "Query does not include newly added employee with null manager_id."
        )
        # Clean up
        self.cursor.execute("DELETE FROM employees WHERE employee_id = 999")
        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):
    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: Find Employees Without Managers

解答