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_columns_present(self):
        """Check if all required columns are present in the query result."""
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        columns = [desc[0] for desc in self.cursor.description]
        required = ["employee_id", "first_name", "last_name", "salary", "department_id"]
        missing = [col for col in required if col not in columns]
        _dynamic_test(
            self,
            len(missing) == 0,
            "All required columns are present.",
            f"Missing columns: {', '.join(missing)}"
        )

    def test_only_high_earners(self):
        """Check that only employees with salary greater than their department's average are returned."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Get all employees
        self.cursor.execute("SELECT employee_id, salary, department_id FROM employees")
        all_emps = pd.DataFrame(self.cursor.fetchall(), columns=["employee_id", "salary", "department_id"])
        # Calculate dept averages
        dept_avg = all_emps.groupby("department_id")['salary'].mean().to_dict()
        # Get user result
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(rows, columns=columns)
        # Check: all returned employees have salary > avg in their department
        failed = []
        for _, row in result_df.iterrows():
            if not row["salary"] > dept_avg[row["department_id"]]:
                failed.append(row["employee_id"])
        _dynamic_test(
            self,
            len(failed) == 0,
            "All returned employees have salary greater than their department's average.",
            f"The following employees are not high earners in their department: {failed}"
        )

    def test_no_single_employee_departments(self):
        """Check that departments with only one employee are not included in the result."""
        with open("query.sql", "r") as file:
            query = file.read()
        # Get department employee counts
        self.cursor.execute("SELECT department_id, COUNT(*) as emp_count FROM employees GROUP BY department_id")
        dept_counts = dict(self.cursor.fetchall())
        single_emp_depts = [dept for dept, count in dept_counts.items() if count == 1]
        # Get user result
        self.cursor.execute(query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(rows, columns=columns)
        # Check if any result row is from a single-employee department
        bad = result_df[result_df['department_id'].isin(single_emp_depts)]
        _dynamic_test(
            self,
            bad.empty,
            "No employees from single-employee departments are included.",
            f"Employees from single-employee departments are included: {bad['employee_id'].tolist()}"
        )

    def test_correlated_subquery_used(self):
        """Check that a correlated subquery referencing the outer department_id is used in the WHERE clause."""
        with open("query.sql", "r") as file:
            query = file.read().lower().replace('\\n', ' ')
        # Look for a subquery in WHERE referencing department_id from outer query
        condition = (
            "where" in query and
            "select" in query and
            "avg(" in query and
            (
                "department_id = e.department_id" in query or
                "department_id = employees.department_id" in query or
                "department_id = department_id" in query
            )
        )
        _dynamic_test(
            self,
            condition,
            "Correlated subquery referencing department_id is used in WHERE.",
            "Query does not use a correlated subquery referencing the outer department_id in WHERE."
        )


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

Master the art of writing and optimizing nested subqueries in SQL. This course explores advanced querying techniques using subqueries in WHERE, FROM, and SELECT clauses, and teaches best practices for performance optimization. Each section blends engaging theory with real-world challenges to solidify your understanding.

Explore how to use nested subqueries within the WHERE clause to filter data based on complex conditions and relationships.

Learn how to use nested subqueries in the FROM clause to create derived tables for complex data analysis.

Discover how to use nested subqueries in the SELECT clause to compute values for each row in your result set.

Master best practices and techniques for optimizing nested subqueries for performance and maintainability.

Challenge: Find High Earners by Department

Solution