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_bonus_column_exists(self):
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute(query)
        columns = [desc[0] for desc in self.cursor.description]
        _dynamic_test(
            self,
            "bonus" in [col.lower() for col in columns],
            "The output includes a 'bonus' column.",
            "The output does not include a 'bonus' column."
        )

    def test_all_employees_included(self):
        with open("query.sql", "r") as file:
            query = file.read()
        self.cursor.execute("SELECT COUNT(*) FROM employees;")
        expected_count = self.cursor.fetchone()[0]
        self.cursor.execute(query)
        result = self.cursor.fetchall()
        _dynamic_test(
            self,
            len(result) == expected_count,
            "All employees are included in the output.",
            f"Expected {expected_count} rows, got {len(result)}."
        )

    def test_high_bonus_logic(self):
        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)
        # High: salary >= 80000 AND years_at_company >= 5
        mask = (df['salary'] >= 80000) & (df['years_at_company'] >= 5)
        high_bonus = df[mask]
        failed = high_bonus[high_bonus['bonus'].astype(str).str.lower() != 'high']
        condition = len(failed) == 0
        if not condition:
            failure_msg = f"Rows with salary >=80000 and years_at_company >=5 do not have 'High' bonus: {failed[['id','bonus']].to_dict('records')}"
        else:
            failure_msg = ""
        _dynamic_test(
            self,
            condition,
            "'High' bonus is correctly assigned.",
            failure_msg
        )

    def test_medium_bonus_logic(self):
        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)
        # Medium: (salary >= 80000 XOR years_at_company >= 5)
        mask_salary = (df['salary'] >= 80000)
        mask_years = (df['years_at_company'] >= 5)
        mask_medium = (mask_salary ^ mask_years)
        medium_bonus = df[mask_medium]
        failed = medium_bonus[medium_bonus['bonus'].astype(str).str.lower() != 'medium']
        condition = len(failed) == 0
        if not condition:
            failure_msg = f"Rows with only one of salary >=80000 or years_at_company >=5 do not have 'Medium' bonus: {failed[['id','bonus']].to_dict('records')}"
        else:
            failure_msg = ""
        _dynamic_test(
            self,
            condition,
            "'Medium' bonus is correctly assigned.",
            failure_msg
        )

    def test_none_bonus_logic(self):
        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)
        # None: salary < 80000 AND years_at_company < 5
        mask = (df['salary'] < 80000) & (df['years_at_company'] < 5)
        none_bonus = df[mask]
        failed = none_bonus[none_bonus['bonus'].astype(str).str.lower() != 'none']
        condition = len(failed) == 0
        if not condition:
            failure_msg = f"Rows with salary <80000 and years_at_company <5 do not have 'None' bonus: {failed[['id','bonus']].to_dict('records')}"
        else:
            failure_msg = ""
        _dynamic_test(
            self,
            condition,
            "'None' bonus is correctly assigned.",
            failure_msg
        )

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 module designed to help you master the CASE WHEN statement in SQL, from basic conditional logic to advanced use cases in data analysis and reporting. Through engaging theory and real-world tasks, you'll learn to write expressive, powerful queries that adapt to complex business requirements.

Learn the basics of the CASE WHEN statement, its syntax, and how to use it for simple conditional logic in SQL queries.

Challenge: Employee Bonus Eligibility

Рішення