import unittest
import psycopg2
import pandas as pd
import json

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()
        # Setup test data
        cls.cursor.execute("DROP TABLE IF EXISTS user_profiles;")
        cls.cursor.execute("""
            CREATE TABLE user_profiles (
                user_id SERIAL PRIMARY KEY,
                profile_data JSONB
            );
        """)
        # Insert test data
        test_data = [
            {"address": {"city": "Springfield"}},
            {"address": {"city": "Greenville"}},
            {"address": {"city": "Fairview"}},
            {"address": {"city": "Lakeside"}},
            {"address": {"city": "Hilltown"}}
        ]
        for d in test_data:
            cls.cursor.execute("INSERT INTO user_profiles (profile_data) VALUES (%s)", [json.dumps(d)])
        cls.connection.commit()

    @classmethod
    def tearDownClass(cls):
        try:
            cls.cursor.execute("ROLLBACK;")
        except Exception:
            pass
        cls.cursor.execute("DROP TABLE IF EXISTS user_profiles;")
        cls.connection.commit()
        cls.cursor.close()
        cls.connection.close()

    def test_returns_springfield_user(self):
        """Check that the user with address.city 'Springfield' is returned."""
        with open("query.sql", "r") as file:
            query = file.read()
        try:
            self.cursor.execute("ROLLBACK;")
        except Exception:
            pass
        try:
            self.cursor.execute(query)
            rows = self.cursor.fetchall()
            # Only one user should be returned
            condition = len(rows) == 1
            failure_msg = f"Expected 1 user, got {len(rows)}: {rows}"
            _dynamic_test(
                self,
                condition,
                "User with address.city 'Springfield' is returned.",
                failure_msg
            )
            # The returned user's city should be Springfield
            if rows:
                profile_data = rows[0][1]
                if isinstance(profile_data, dict):
                    city = profile_data.get('address',{}).get('city')
                else:
                    d = json.loads(profile_data)
                    city = d.get('address',{}).get('city')
                _dynamic_test(
                    self,
                    city == "Springfield",
                    "Returned user's city is Springfield.",
                    f"Returned user's city is '{city}', expected 'Springfield'."
                )
        except Exception as e:
            _dynamic_test(
                self,
                False,
                "User with address.city 'Springfield' is returned.",
                f"Query execution failed: {e}"
            )

    def test_returns_columns(self):
        """Check that the query returns both user_id and profile_data columns."""
        with open("query.sql", "r") as file:
            query = file.read()
        try:
            self.cursor.execute("ROLLBACK;")
        except Exception:
            pass
        try:
            self.cursor.execute(query)
            columns = [desc[0] for desc in self.cursor.description]
            expected = set(["user_id", "profile_data"])
            _dynamic_test(
                self,
                expected.issubset(set(columns)),
                "Both user_id and profile_data columns are present.",
                f"Returned columns are missing expected columns: {columns}"
            )
        except Exception as e:
            _dynamic_test(
                self,
                False,
                "Both user_id and profile_data columns are present.",
                f"Query execution failed: {e}"
            )

    def test_no_false_positives(self):
        """Ensure query does not return users whose address.city is not 'Springfield'."""
        with open("query.sql", "r") as file:
            query = file.read()
        try:
            self.cursor.execute("ROLLBACK;")
        except Exception:
            pass
        try:
            self.cursor.execute(query)
            rows = self.cursor.fetchall()
            for row in rows:
                profile_data = row[1]
                if isinstance(profile_data, dict):
                    city = profile_data.get('address',{}).get('city')
                else:
                    d = json.loads(profile_data)
                    city = d.get('address',{}).get('city')
                _dynamic_test(
                    self,
                    city == "Springfield",
                    "No users with city other than 'Springfield' returned.",
                    f"Returned user with city '{city}', but only 'Springfield' should be returned."
                )
        except Exception as e:
            _dynamic_test(
                self,
                False,
                "No users with city other than 'Springfield' returned.",
                f"Query execution failed: {e}"
            )

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):
    import re
    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 Users by City

解答