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()
        # Drop and create table
        cls.cursor.execute("DROP TABLE IF EXISTS country;")
        cls.cursor.execute("""
            CREATE TABLE country (
                id SERIAL PRIMARY KEY,
                name VARCHAR(100),
                population INTEGER,
                region VARCHAR(100)
            );
        """)
        # Insert sample data
        sample_data = [
            (1, 'France', 67000000, 'Europe'),
            (2, 'Brazil', 210000000, 'South America'),
            (3, 'Canada', 37000000, 'North America'),
            (4, 'Japan', 126000000, 'Asia'),
            (5, 'Nigeria', 200000000, 'Africa'),
            (6, 'Germany', 83000000, 'Europe'),
            (7, 'Chile', 18000000, 'South America'),
            (8, 'Egypt', 100000000, 'Africa'),
            (9, 'China', 1393000000, 'Asia'),
            (10, 'Argentina', 45000000, 'South America')
        ]
        cls.cursor.execute("DELETE FROM country;")
        for row in sample_data:
            cls.cursor.execute(
                "INSERT INTO country (id, name, population, region) VALUES (%s, %s, %s, %s)",
                row
            )
        cls.connection.commit()

    def test_selects_correct_columns(self):
        """Check if query selects the correct columns in the correct order."""
        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 = ["id", "name", "population", "region"]
        condition = columns == expected_columns
        _dynamic_test(
            self,
            condition,
            "Query selects the correct columns in the correct order.",
            f"Query columns are {columns}, but expected {expected_columns}."
        )

    def test_orders_by_region_desc_population_asc(self):
        """Check if query orders by region descending and then by population ascending."""
        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)
        if df.empty:
            _dynamic_test(
                self,
                True,
                "Query orders by region descending and population ascending (no rows returned).",
                "Query did not return any rows."
            )
            return
        expected_df = df.sort_values(by=["region", "population"], ascending=[False, True], na_position="last").reset_index(drop=True)
        condition = df.reset_index(drop=True).equals(expected_df)
        _dynamic_test(
            self,
            condition,
            "Query orders by region descending and population ascending.",
            "Query does not order by region descending and population ascending."
        )

    def test_returns_expected_rows(self):
        """Check if query returns the expected rows and values based on data."""
        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)
        # Build expected result using the correct query
        self.cursor.execute("""
            SELECT id, name, population, region  
            FROM country 
            ORDER BY region DESC, population;
        """)
        expected_rows = self.cursor.fetchall()
        expected_columns = [desc[0] for desc in self.cursor.description]
        expected_df = pd.DataFrame(expected_rows, columns=expected_columns)
        # Compare all rows and values
        condition = result_df.equals(expected_df)
        _dynamic_test(
            self,
            condition,
            "Query returns the expected rows and values in order.",
            "Query does not return the expected rows and values in order."
        )

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

Gain a solid understanding of SQL essentials tailored specifically for business intelligence. You will master querying, filtering, and aggregating business datasets through hands-on tasks designed to mirror real-world professional scenarios.

Challenge: Find Countries, IDs, Populations, Regions, and Sort Them

Рішення