Class PGTest

compatibility/pg/python/pg_test.py:3–93

import psycopg

class PGTest:
    class Test:
        def __init__(self, query, expected_results):
            self.query = query
            self.expected_results = expected_results

        def run(self, cursor):
            print(f"Running test: {self.query}")
            cursor.execute(self.query)
            if cursor.description is None:
                if len(self.expected_results) != 0:
                    print(f"Expected {len(self.expected_results)} rows, got 0")
                    return False
                print("Returns 0 rows")
                return True

            rows = cursor.fetchall()
            if len(rows[0]) != len(self.expected_results[0]):
                print(f"Expected {len(self.expected_results[0])} columns, got {len(rows[0])}")
                return False

            for row, expected_row in zip(rows, self.expected_results):
                if list(map(str, row)) != list(map(str, expected_row)):
                    print(f"Expected: {list(map(str, expected_row))}, got: {list(map(str, row))}")
                    return False

            print(f"Returns {len(rows)} rows")
            if len(rows) != len(self.expected_results):
                print(f"Expected {len(self.expected_results)} rows")
                return False

            return True

    def __init__(self):
        self.conn = None
        self.tests = []

    def connect(self, ip, port, user, password):
        try:
            self.conn = psycopg.connect(
                host=ip,
                port=port,
                dbname="postgres",
                user=user,
                password=password,
                autocommit=True
            )
        except Exception as e:
            raise RuntimeError(e)

    def disconnect(self):
        self.conn.close()

    def add_test(self, query, expected_results):
        self.tests.append(self.Test(query, expected_results))

    def run_tests(self):
        for test in self.tests: