# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
import json
import re

from age.models import Vertex, Edge, Path
import unittest
import unittest.mock
import decimal
import age
# _validate_column is private but tested directly because its quoting
# behavior is security-relevant and the public surface (buildCypher)
# makes it difficult to isolate quoting assertions.
from age.age import buildCypher, _validate_column
from age.exceptions import InvalidIdentifier
import argparse

TEST_HOST = "localhost"
TEST_PORT = 5432
TEST_DB = "postgres"
TEST_USER = "postgres"
TEST_PASSWORD = "agens"
TEST_GRAPH_NAME = "test_graph"


class TestSetUpAge(unittest.TestCase):
    """Unit tests for setUpAge() skip_load parameter — no DB required."""

    def _make_mock_conn(self):
        mock_conn = unittest.mock.MagicMock()
        mock_cursor = unittest.mock.MagicMock()
        mock_conn.cursor.return_value.__enter__ = unittest.mock.Mock(return_value=mock_cursor)
        mock_conn.cursor.return_value.__exit__ = unittest.mock.Mock(return_value=False)
        mock_conn.adapters = unittest.mock.MagicMock()
        mock_type_info = unittest.mock.MagicMock()
        mock_type_info.oid = 1
        mock_type_info.array_oid = 2
        return mock_conn, mock_cursor, mock_type_info

    def test_skip_load_true_does_not_execute_load(self):
        """When skip_load=True, LOAD 'age' must not be executed."""
        mock_conn, mock_cursor, mock_type_info = self._make_mock_conn()
        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=mock_type_info), \
             unittest.mock.patch("age.age.checkGraphCreated"):
            age.age.setUpAge(mock_conn, "test_graph", skip_load=True)
        mock_cursor.execute.assert_called_once_with(
            'SET search_path = ag_catalog, "$user", public;'
        )

    def test_skip_load_false_executes_load(self):
        """When skip_load=False (default), LOAD 'age' must be executed."""
        mock_conn, mock_cursor, mock_type_info = self._make_mock_conn()
        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=mock_type_info), \
             unittest.mock.patch("age.age.checkGraphCreated"):
            age.age.setUpAge(mock_conn, "test_graph", skip_load=False)
        mock_cursor.execute.assert_any_call("LOAD 'age';")

    def test_skip_load_with_load_from_plugins(self):
        """When skip_load=False and load_from_plugins=True, LOAD from plugins path."""
        mock_conn, mock_cursor, mock_type_info = self._make_mock_conn()
        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=mock_type_info), \
             unittest.mock.patch("age.age.checkGraphCreated"):
            age.age.setUpAge(mock_conn, "test_graph", load_from_plugins=True, skip_load=False)
        mock_cursor.execute.assert_any_call("LOAD '$libdir/plugins/age';")

    def test_skip_load_true_still_sets_search_path(self):
        """When skip_load=True, search_path must still be set."""
        mock_conn, mock_cursor, mock_type_info = self._make_mock_conn()
        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=mock_type_info), \
             unittest.mock.patch("age.age.checkGraphCreated"):
            age.age.setUpAge(mock_conn, "test_graph", skip_load=True)
        mock_cursor.execute.assert_any_call(
            'SET search_path = ag_catalog, "$user", public;'
        )

    def test_contradictory_skip_load_and_load_from_plugins_raises(self):
        """skip_load=True + load_from_plugins=True must raise ValueError."""
        mock_conn, _, _ = self._make_mock_conn()
        with self.assertRaises(ValueError):
            age.age.setUpAge(mock_conn, "test_graph", load_from_plugins=True, skip_load=True)

    def test_connect_forwards_skip_load_to_setup(self):
        """age.connect(skip_load=True) must forward skip_load through the full call chain."""
        with unittest.mock.patch("age.age.psycopg.connect") as mock_psycopg, \
             unittest.mock.patch("age.age.setUpAge") as mock_setup:
            mock_psycopg.return_value = unittest.mock.MagicMock()
            age.connect(dsn="host=localhost", graph="test_graph", skip_load=True)
        mock_setup.assert_called_once()
        _, kwargs = mock_setup.call_args
        self.assertTrue(
            kwargs.get("skip_load", False),
            "skip_load must be forwarded from age.connect() to setUpAge()"
        )


class TestBuildCypher(unittest.TestCase):
    """Unit tests for buildCypher() and _validate_column() — no DB required."""

    def test_simple_column(self):
        result = buildCypher("g", "MATCH (n) RETURN n", ["n"])
        self.assertIn('"n" agtype', result)

    def test_column_with_type(self):
        result = buildCypher("g", "MATCH (n) RETURN n", ["n agtype"])
        self.assertIn('"n" agtype', result)

    def test_reserved_word_count(self):
        """Issue #2370: 'count' is a PostgreSQL reserved word."""
        result = buildCypher("g", "MATCH (n) RETURN count(n)", ["count"])
        self.assertIn('"count" agtype', result)
        # Verify 'count' never appears unquoted as a column name
        self.assertIsNone(
            re.search(r'(?<!")\bcount\s+agtype\b', result),
            f"'count' must be quoted in: {result}"
        )

    def test_reserved_word_order(self):
        """Issue #2370: 'order' is a PostgreSQL reserved word."""
        result = buildCypher("g", "MATCH (n) RETURN n.order", ["order"])
        self.assertIn('"order" agtype', result)

    def test_reserved_word_type(self):
        """Issue #2370: 'type' is a PostgreSQL reserved word."""
        result = buildCypher("g", "MATCH ()-[r]->() RETURN type(r)", ["type"])
        self.assertIn('"type" agtype', result)

    def test_reserved_word_select(self):
        """Issue #2370: 'select' is a PostgreSQL reserved word."""
        result = buildCypher("g", "MATCH (n) RETURN n", ["select"])
        self.assertIn('"select" agtype', result)

    def test_reserved_word_group(self):
        """Issue #2370: 'group' is a PostgreSQL reserved word."""
        result = buildCypher("g", "MATCH (n) RETURN n", ["group"])
        self.assertIn('"group" agtype', result)

    def test_multiple_columns(self):
        result = buildCypher("g", "MATCH (n) RETURN n.name, count(n)", ["name", "count"])
        self.assertIn('"name" agtype', result)
        self.assertIn('"count" agtype', result)

    def test_default_column(self):
        result = buildCypher("g", "MATCH (n) RETURN n", None)
        self.assertIn('"v" agtype', result)

    def test_invalid_column_rejected(self):
        with self.assertRaises(InvalidIdentifier):
            buildCypher("g", "MATCH (n) RETURN n", ["invalid;col"])

    def test_reserved_word_in_name_type_pair(self):
        """Quoting applies even when the column is specified as 'name type'."""
        result = buildCypher("g", "MATCH (n) RETURN n.order", ["order agtype"])
        self.assertIn('"order" agtype', result)

    def test_validate_column_quoting(self):
        self.assertEqual(_validate_column("v"), '"v" agtype')
        self.assertEqual(_validate_column("v agtype"), '"v" agtype')
        self.assertEqual(_validate_column("count"), '"count" agtype')
        self.assertEqual(_validate_column("my_col"), '"my_col" agtype')


class TestModelToDict(unittest.TestCase):
    """Unit tests for Vertex/Edge/Path to_dict() — no DB required."""

    def test_vertex_to_dict(self):
        v = Vertex(id=123, label="Person", properties={"name": "Alice", "age": 30})
        d = v.to_dict()
        self.assertEqual(d["id"], 123)
        self.assertEqual(d["label"], "Person")
        self.assertEqual(d["properties"], {"name": "Alice", "age": 30})
        # Verify it's a plain dict (JSON-serializable)
        json_str = json.dumps(d)
        self.assertIn("Alice", json_str)

    def test_vertex_to_dict_empty_properties(self):
        v = Vertex(id=1, label="Empty", properties=None)
        d = v.to_dict()
        self.assertEqual(d["properties"], {})

    def test_edge_to_dict(self):
        e = Edge(id=456, label="KNOWS", properties={"since": 2020})
        e.start_id = 123
        e.end_id = 789
        d = e.to_dict()
        self.assertEqual(d["id"], 456)
        self.assertEqual(d["label"], "KNOWS")
        self.assertEqual(d["start_id"], 123)
        self.assertEqual(d["end_id"], 789)
        self.assertEqual(d["properties"], {"since": 2020})
        json_str = json.dumps(d)
        self.assertIn("KNOWS", json_str)

    def test_path_to_dict(self):
        v1 = Vertex(id=1, label="A", properties={"name": "start"})
        e = Edge(id=10, label="r", properties={"w": 1})
        e.start_id = 1
        e.end_id = 2
        v2 = Vertex(id=2, label="B", properties={"name": "end"})
        p = Path([v1, e, v2])
        d = p.to_dict()
        self.assertEqual(len(d), 3)
        self.assertEqual(d[0]["label"], "A")
        self.assertEqual(d[1]["label"], "r")
        self.assertEqual(d[1]["start_id"], 1)
        self.assertEqual(d[2]["label"], "B")
        # Verify the whole path is JSON-serializable
        json_str = json.dumps(d)
        self.assertIn("start", json_str)

    def test_vertex_to_dict_is_plain_dict(self):
        """to_dict() returns standard dict, not a model object."""
        v = Vertex(id=1, label="X", properties={"k": "v"})
        d = v.to_dict()
        self.assertIsInstance(d, dict)
        self.assertIsInstance(d["properties"], dict)


class TestPublicImports(unittest.TestCase):
    """Verify that public API symbols are importable without type: ignore."""

    def test_import_configure_connection(self):
        from age import configure_connection
        self.assertTrue(callable(configure_connection))

    def test_import_age_loader(self):
        from age import AgeLoader
        self.assertIsNotNone(AgeLoader)

    def test_import_client_cursor(self):
        from age import ClientCursor
        self.assertIsNotNone(ClientCursor)


class TestConfigureConnection(unittest.TestCase):
    """Unit tests for configure_connection() — no DB required."""

    def _make_mock_conn(self):
        mock_conn = unittest.mock.MagicMock()
        mock_cursor = unittest.mock.MagicMock()
        mock_conn.cursor.return_value.__enter__ = unittest.mock.Mock(return_value=mock_cursor)
        mock_conn.cursor.return_value.__exit__ = unittest.mock.Mock(return_value=False)
        mock_conn.adapters = unittest.mock.MagicMock()
        mock_type_info = unittest.mock.MagicMock()
        mock_type_info.oid = 1
        mock_type_info.array_oid = 2
        return mock_conn, mock_cursor, mock_type_info

    def test_default_does_not_load(self):
        """By default, configure_connection should NOT execute LOAD."""
        mock_conn, mock_cursor, mock_type_info = self._make_mock_conn()
        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=mock_type_info), \
             unittest.mock.patch("age.age.checkGraphCreated"):
            age.age.configure_connection(mock_conn)
        mock_cursor.execute.assert_called_once_with(
            'SET search_path = ag_catalog, "$user", public;'
        )

    def test_load_true_executes_load(self):
        """When load=True, LOAD 'age' must be executed."""
        mock_conn, mock_cursor, mock_type_info = self._make_mock_conn()
        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=mock_type_info), \
             unittest.mock.patch("age.age.checkGraphCreated"):
            age.age.configure_connection(mock_conn, load=True)
        mock_cursor.execute.assert_any_call("LOAD 'age';")

    def test_load_from_plugins(self):
        """When load=True and load_from_plugins=True, use plugins path."""
        mock_conn, mock_cursor, mock_type_info = self._make_mock_conn()
        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=mock_type_info), \
             unittest.mock.patch("age.age.checkGraphCreated"):
            age.age.configure_connection(mock_conn, load=True, load_from_plugins=True)
        mock_cursor.execute.assert_any_call("LOAD '$libdir/plugins/age';")

    def test_load_from_plugins_without_load_raises(self):
        """load_from_plugins=True without load=True must raise ValueError."""
        mock_conn, _, _ = self._make_mock_conn()
        with self.assertRaises(ValueError):
            age.age.configure_connection(mock_conn, load_from_plugins=True)

    def test_always_sets_search_path(self):
        """search_path must always be set regardless of load parameter."""
        mock_conn, mock_cursor, mock_type_info = self._make_mock_conn()
        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=mock_type_info), \
             unittest.mock.patch("age.age.checkGraphCreated"):
            age.age.configure_connection(mock_conn)
        mock_cursor.execute.assert_any_call(
            'SET search_path = ag_catalog, "$user", public;'
        )

    def test_registers_agtype_adapters(self):
        """AgeLoader must be registered for agtype OIDs."""
        mock_conn, mock_cursor, mock_type_info = self._make_mock_conn()
        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=mock_type_info), \
             unittest.mock.patch("age.age.checkGraphCreated"):
            age.age.configure_connection(mock_conn)
        mock_conn.adapters.register_loader.assert_any_call(1, age.age.AgeLoader)
        mock_conn.adapters.register_loader.assert_any_call(2, age.age.AgeLoader)

    def test_graph_name_triggers_check(self):
        """When graph_name is provided, checkGraphCreated must be called."""
        mock_conn, mock_cursor, mock_type_info = self._make_mock_conn()
        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=mock_type_info), \
             unittest.mock.patch("age.age.checkGraphCreated") as mock_check:
            age.age.configure_connection(mock_conn, graph_name="my_graph")
        mock_check.assert_called_once_with(mock_conn, "my_graph")

    def test_age_not_set_when_type_info_is_none(self):
        """AgeNotSet must be raised when TypeInfo.fetch returns None."""
        from age.exceptions import AgeNotSet
        mock_conn, _, _ = self._make_mock_conn()
        with unittest.mock.patch("age.age.TypeInfo.fetch", return_value=None):
            with self.assertRaises(AgeNotSet):
                age.age.configure_connection(mock_conn)


class TestAgeBasic(unittest.TestCase):
    ag = None
    args: argparse.Namespace = argparse.Namespace(
        host=TEST_HOST,
        port=TEST_PORT,
        database=TEST_DB,
        user=TEST_USER,
        password=TEST_PASSWORD,
        graphName=TEST_GRAPH_NAME
    )

    def setUp(self):
        print("Connecting to Test Graph.....")
        args = dict(
            host=self.args.host,
            port=self.args.port,
            dbname=self.args.database,
            user=self.args.user,
            password=self.args.password,
        )

        dsn = "host={host} port={port} dbname={dbname} user={user} password={password}".format(
            **args
        )
        self.ag = age.connect(dsn, graph=self.args.graphName, **args)

    def tearDown(self):
        # Clear test data
        print("Deleting Test Graph.....")
        age.deleteGraph(self.ag.connection, self.ag.graphName)
        self.ag.close()

    def testExec(self):
        print("\n---------------------------------------------------")
        print("Test 1: Checking single and multi column Returns.....")
        print("---------------------------------------------------\n")

        ag = self.ag
        # Create and Return single column
        cursor = ag.execCypher(
            "CREATE (n:Person {name: %s, title: 'Developer'}) RETURN n",
            params=("Andy",),
        )
        for row in cursor:
            print("Vertex: %s , Type: %s " % (Vertex, type(row[0])))

        # Create and Return multi columns
        cursor = ag.execCypher(
            "CREATE (n:Person {name: %s, title: %s}) RETURN id(n), n.name",
            cols=["id", "name"],
            params=("Jack", "Manager"),
        )
        row = cursor.fetchone()
        print("Id: %s , Name: %s" % (row[0], row[1]))
        self.assertEqual(int, type(row[0]))
        ag.commit()
        print("\nTest 1 Successful....")

    def testQuery(self):
        print("\n--------------------------------------------------")
        print("Test 2: Testing CREATE and query relationships.....")
        print("--------------------------------------------------\n")

        ag = self.ag
        ag.execCypher("CREATE (n:Person {name: %s}) ", params=("Jack",))
        ag.execCypher("CREATE (n:Person {name: %s}) ", params=("Andy",))
        ag.execCypher("CREATE (n:Person {name: %s}) ", params=("Smith",))
        ag.execCypher(
            "MATCH (a:Person), (b:Person) WHERE a.name = 'Andy' AND b.name = 'Jack' CREATE (a)-[r:worksWith {weight: 3}]->(b)"
        )
        ag.execCypher(
            """MATCH (a:Person), (b:Person)
                    WHERE  a.name = %s AND b.name = %s
                    CREATE p=((a)-[r:worksWith]->(b)) """,
            params=(
                "Jack",
                "Smith",
            ),
        )

        ag.commit()

        cursor = ag.execCypher("MATCH p=()-[:worksWith]-() RETURN p")
        for row in cursor:
            path = row[0]
            print("START:", path[0])
            print("EDGE:", path[1])
            print("END:", path[2])

        cursor = ag.execCypher(
            "MATCH p=(a)-[b]-(c) WHERE b.weight>%s RETURN a,label(b), b.weight, c",
            cols=["a", "bl", "bw", "c"],
            params=(2,),
        )
        for row in cursor:
            start = row[0]
            edgel = row[1]
            edgew = row[2]
            end = row[3]
            print(
                "Relationship: %s %s %s. Edge weight: %s"
                % (start["name"], edgel, end["name"], edgew)
            )
            # Assert that the weight of the edge is greater than 2
            self.assertEqual(edgew > 2, True)
        print("\nTest 2 Successful...")

    def testChangeData(self):
        print("\n-------------------------------------------------------")
        print("Test 3: Testing changes in data using SET and REMOVE.....")
        print("-------------------------------------------------------\n")

        ag = self.ag
        # Create Vertices
        # Commit automatically
        ag.execCypher("CREATE (n:Person {name: 'Joe'})")

        cursor = ag.execCypher(
            "CREATE (n:Person {name: %s, title: 'Developer'}) RETURN n",
            params=("Smith",),
        )
        row = cursor.fetchone()
        print("CREATED: ", row[0])

        # You must commit explicitly
        ag.commit()

        cursor = ag.execCypher(
            "MATCH (n:Person {name: %s}) SET n.title=%s RETURN n",
            params=(
                "Smith",
                "Manager",
            ),
        )
        row = cursor.fetchone()
        vertex = row[0]
        title1 = vertex["title"]
        print("SET title: ", title1)

        ag.commit()

        cursor = ag.execCypher("MATCH (p:Person {name: 'Smith'}) RETURN p.title")
        row = cursor.fetchone()
        title2 = row[0]

        self.assertEqual(title1, title2)

        cursor = ag.execCypher(
            "MATCH (n:Person {name: %s}) SET n.bigNum=-6.45161e+46::numeric RETURN n",
            params=("Smith",),
        )
        row = cursor.fetchone()
        vertex = row[0]
        for row in cursor:
            print("SET bigNum: ", vertex["bigNum"])

        bigNum1 = vertex["bigNum"]

        self.assertEqual(decimal.Decimal("-6.45161e+46"), bigNum1)
        ag.commit()

        cursor = ag.execCypher("MATCH (p:Person {name: 'Smith'}) RETURN p.bigNum")
        row = cursor.fetchone()
        bigNum2 = row[0]

        self.assertEqual(bigNum1, bigNum2)

        cursor = ag.execCypher(
            "MATCH (n:Person {name: %s}) REMOVE n.title RETURN n", params=("Smith",)
        )
        for row in cursor:
            print("REMOVE Prop title: ", row[0])
            # Assert that the title property is removed
            self.assertIsNone(row[0].properties.get("title"))
        print("\nTest 3 Successful....")

        # You must commit explicitly
        ag.commit()

    def testCypher(self):
        print("\n--------------------------")
        print("Test 4: Testing Cypher.....")
        print("--------------------------\n")

        ag = self.ag

        with ag.connection.cursor() as cursor:
            try:
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Joe",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Jack",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Andy",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Smith",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Tom",))

                # You must commit explicitly
                ag.commit()
            except Exception as ex:
                print(ex)
                ag.rollback()

        with ag.connection.cursor() as cursor:
            try:  # Create Edges
                ag.cypher(
                    cursor,
                    "MATCH (a:Person), (b:Person) WHERE a.name = 'Joe' AND b.name = 'Smith' CREATE (a)-[r:worksWith {weight: 3}]->(b)",
                )
                ag.cypher(
                    cursor,
                    "MATCH (a:Person), (b:Person) WHERE  a.name = 'Andy' AND b.name = 'Tom' CREATE (a)-[r:worksWith {weight: 1}]->(b)",
                )
                ag.cypher(
                    cursor,
                    "MATCH (a:Person {name: 'Jack'}), (b:Person {name: 'Andy'}) CREATE (a)-[r:worksWith {weight: 5}]->(b)",
                )

                # You must commit explicitly
                ag.commit()
            except Exception as ex:
                print(ex)
                ag.rollback()

        # With Params
        cursor = ag.execCypher(
            """MATCH (a:Person), (b:Person)
                WHERE  a.name = %s AND b.name = %s
                CREATE p=((a)-[r:worksWith]->(b)) RETURN p""",
            params=(
                "Andy",
                "Smith",
            ),
        )

        for row in cursor:
            print("CREATED EDGE: %s" % row[0])

        cursor = ag.execCypher("""MATCH (a:Person {name: 'Joe'}), (b:Person {name: 'Jack'})
                CREATE p=((a)-[r:worksWith {weight: 5}]->(b))
                RETURN p """)

        for row in cursor:
            print("CREATED EDGE WITH PROPERTIES: %s" % row[0])
            self.assertEqual(row[0][1].properties["weight"], 5)

        print("\nTest 4 Successful...")

    def testMultipleEdges(self):
        print("\n------------------------------------")
        print("Test 5: Testing Multiple Edges.....")
        print("------------------------------------\n")

        ag = self.ag
        with ag.connection.cursor() as cursor:
            try:
                ag.cypher(cursor, "CREATE (n:Country {name: %s}) ", params=("USA",))
                ag.cypher(cursor, "CREATE (n:Country {name: %s}) ", params=("France",))
                ag.cypher(cursor, "CREATE (n:Country {name: %s}) ", params=("Korea",))
                ag.cypher(cursor, "CREATE (n:Country {name: %s}) ", params=("Russia",))

                # You must commit explicitly after all executions.
                ag.connection.commit()
            except Exception as ex:
                ag.rollback()
                raise ex

        with ag.connection.cursor() as cursor:
            try:  # Create Edges
                ag.cypher(
                    cursor,
                    "MATCH (a:Country), (b:Country) WHERE a.name = 'USA' AND b.name = 'France' CREATE (a)-[r:distance {unit:'miles', value: 4760}]->(b)",
                )
                ag.cypher(
                    cursor,
                    "MATCH (a:Country), (b:Country) WHERE  a.name = 'France' AND b.name = 'Korea' CREATE (a)-[r:distance {unit: 'km', value: 9228}]->(b)",
                )
                ag.cypher(
                    cursor,
                    "MATCH (a:Country {name: 'Korea'}), (b:Country {name: 'Russia'}) CREATE (a)-[r:distance {unit:'km', value: 3078}]->(b)",
                )

                # You must commit explicitly
                ag.connection.commit()
            except Exception as ex:
                ag.rollback()
                raise ex

        cursor = ag.execCypher("""MATCH p=(:Country {name:"USA"})-[:distance]-(:Country)-[:distance]-(:Country)
                RETURN p""")

        count = 0
        output = []
        for row in cursor:
            path = row[0]
            for e in path:
                if e.gtype == age.TP_VERTEX:
                    output.append(e.label + " " + e["name"])
                elif e.gtype == age.TP_EDGE:
                    output.append(
                        "---- (distance " + str(e["value"]) + " " + e["unit"] + ") --->"
                    )
                else:
                    output.append("Unknown element. " + str(e))

                count += 1

        formatted_output = " ".join(output)
        print("PATH WITH MULTIPLE EDGES: %s" % formatted_output)
        self.assertEqual(5, count)

        print("\nTest 5 Successful...")

    def testCollect(self):
        print("\n--------------------------")
        print("Test 6: Testing COLLECT.....")
        print("--------------------------\n")

        ag = self.ag

        with ag.connection.cursor() as cursor:
            try:
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Joe",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Jack",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Andy",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Smith",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Tom",))

                # You must commit explicitly
                ag.commit()
            except Exception as ex:
                print(ex)
                ag.rollback()

        with ag.connection.cursor() as cursor:
            try:  # Create Edges
                ag.cypher(
                    cursor,
                    "MATCH (a:Person), (b:Person) WHERE a.name = 'Joe' AND b.name = 'Smith' CREATE (a)-[r:worksWith {weight: 3}]->(b)",
                )
                ag.cypher(
                    cursor,
                    "MATCH (a:Person), (b:Person) WHERE  a.name = 'Joe' AND b.name = 'Tom' CREATE (a)-[r:worksWith {weight: 1}]->(b)",
                )
                ag.cypher(
                    cursor,
                    "MATCH (a:Person {name: 'Joe'}), (b:Person {name: 'Andy'}) CREATE (a)-[r:worksWith {weight: 5}]->(b)",
                )

                # You must commit explicitly
                ag.commit()
            except Exception as ex:
                print(ex)
                ag.rollback()

        print(" -------- TESTING COLLECT #1 --------")
        with ag.connection.cursor() as cursor:
            ag.cypher(
                cursor,
                "MATCH (a)-[:worksWith]->(c) WITH a as V, COLLECT(c) as CV RETURN V.name, CV",
                cols=["V", "CV"],
            )
            for row in cursor:
                nm = row[0]
                collected = row[1]
                print(nm, "worksWith", [i["name"] for i in collected])
                self.assertEqual(3, len(collected))

        print(" -------- TESTING COLLECT #2 --------")
        for row in ag.execCypher(
            "MATCH (a)-[:worksWith]->(c) WITH a as V, COLLECT(c) as CV RETURN V.name, CV",
            cols=["V1", "CV"],
        ):
            nm = row[0]
            collected = row[1]
            print(nm, "worksWith", [i["name"] for i in collected])
            self.assertEqual(3, len(collected))
        print("\nTest 6 Successful...")

    def testSerialization(self):
        print("\n---------------------------------------")
        print("Test 6: Testing Vertex Serialization.....")
        print("-----------------------------------------\n")

        ag = self.ag

        with ag.connection.cursor() as cursor:
            try:
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Joe",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Jack",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Andy",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Smith",))
                ag.cypher(cursor, "CREATE (n:Person {name: %s}) ", params=("Tom",))

                # You must commit explicitly
                ag.commit()
            except Exception as ex:
                print(ex)
                ag.rollback()

        print(" -------- TESTING Output #1 --------")
        cursor = ag.execCypher("MATCH (n) RETURN n")

        for row in cursor:
            vertex = row[0]
            try:
                # json.loads will fail if the json str is not properly formatted
                as_dict = json.loads(vertex.toJson())
                print("Vertex.toJson() returns a correct json string.")
                assert True
            except:
                assert False

        print(" -------- TESTING Output #2 --------")
        cursor = ag.execCypher("MATCH (n) RETURN n")

        for row in cursor:
            vertex = row[0]
            as_str = vertex.toString()
            # Checking if the trailing comma appears in .toString() output
            self.assertFalse(as_str.endswith(", }}::VERTEX"))
        print("Vertex.toString() 'properties' field is formatted properly.")

    def testConfigureConnection(self):
        """Integration: configure_connection() on an externally-opened
        connection must register agtype adapters so cypher queries return
        real Vertex/Edge/Path objects, and to_dict() must round-trip those
        parser-produced objects through json.dumps()."""
        print("\n-------------------------------------------------------")
        print("Test 8: configure_connection + to_dict end-to-end.....")
        print("-------------------------------------------------------\n")

        import psycopg

        from age import configure_connection

        dsn = "host={host} port={port} dbname={database} user={user} password={password}".format(
            **vars(self.args)
        )
        # Deliberately bypass age.connect(): the whole point of
        # configure_connection() is to enable AGE on a caller-managed
        # connection (e.g. one obtained from psycopg_pool.ConnectionPool).
        raw = psycopg.connect(dsn)
        try:
            configure_connection(raw, graph_name=self.args.graphName, load=True)

            graph = self.args.graphName
            with raw.cursor() as cur:
                cur.execute(
                    f"SELECT * FROM cypher('{graph}', $$ "
                    "CREATE (a:Person {name: 'Alice'})-[r:KNOWS {since: 2020}]->(b:Person {name: 'Bob'}) "
                    "RETURN a, r, b "
                    "$$) AS (a agtype, r agtype, b agtype);"
                )
                row = cur.fetchone()
                raw.commit()

            v_a, e, v_b = row
            self.assertIsInstance(v_a, Vertex)
            self.assertIsInstance(e, Edge)
            self.assertIsInstance(v_b, Vertex)
            self.assertEqual(v_a["name"], "Alice")
            self.assertEqual(v_b["name"], "Bob")
            self.assertEqual(e["since"], 2020)

            payload = {
                "start": v_a.to_dict(),
                "edge": e.to_dict(),
                "end": v_b.to_dict(),
            }
            serialised = json.loads(json.dumps(payload))
            self.assertEqual(serialised["start"]["label"], "Person")
            self.assertEqual(serialised["edge"]["label"], "KNOWS")
            self.assertEqual(serialised["edge"]["start_id"], v_a.id)
            self.assertEqual(serialised["edge"]["end_id"], v_b.id)
            self.assertEqual(serialised["start"]["properties"]["name"], "Alice")

            with raw.cursor() as cur:
                cur.execute(
                    f"SELECT * FROM cypher('{graph}', $$ "
                    "MATCH p=(:Person)-[:KNOWS]->(:Person) RETURN p "
                    "$$) AS (p agtype);"
                )
                path = cur.fetchone()[0]

            self.assertIsInstance(path, Path)
            path_dict = json.loads(json.dumps(path.to_dict()))
            self.assertEqual(len(path_dict), 3)
            self.assertEqual(path_dict[0]["label"], "Person")
            self.assertEqual(path_dict[1]["label"], "KNOWS")
            self.assertEqual(path_dict[2]["label"], "Person")
            print("\nTest 8 Successful....")
        finally:
            raw.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "-host",
        "--host",
        help='Optional Host Name. Default Host is "127.0.0.1" ',
        default=TEST_HOST,
    )
    parser.add_argument(
        "-port",
        "--port",
        help="Optional Port Number. Default port no is 5432",
        default=TEST_PORT,
    )
    parser.add_argument(
        "-db", "--database", help="Required Database Name", default=TEST_DB
    )
    parser.add_argument(
        "-u", "--user", help="Required Username Name", default=TEST_USER
    )
    parser.add_argument(
        "-pass",
        "--password",
        help="Required Password for authentication",
        default=TEST_PASSWORD,
    )
    parser.add_argument(
        "-gn",
        "--graphName",
        help='Optional Graph Name to be created. Default graphName is "test_graph"',
        default=TEST_GRAPH_NAME,
    )

    args = parser.parse_args()
    suite = unittest.TestSuite()
    # Unit tests (no DB required)
    loader = unittest.TestLoader()
    suite.addTests(loader.loadTestsFromTestCase(TestSetUpAge))
    suite.addTests(loader.loadTestsFromTestCase(TestBuildCypher))
    suite.addTests(loader.loadTestsFromTestCase(TestModelToDict))
    suite.addTests(loader.loadTestsFromTestCase(TestPublicImports))
    suite.addTests(loader.loadTestsFromTestCase(TestConfigureConnection))
    # Integration tests (require DB)
    suite.addTest(TestAgeBasic("testExec"))
    suite.addTest(TestAgeBasic("testQuery"))
    suite.addTest(TestAgeBasic("testChangeData"))
    suite.addTest(TestAgeBasic("testCypher"))
    suite.addTest(TestAgeBasic("testMultipleEdges"))
    suite.addTest(TestAgeBasic("testCollect"))
    suite.addTest(TestAgeBasic("testSerialization"))
    suite.addTest(TestAgeBasic("testConfigureConnection"))
    TestAgeBasic.args = args
    unittest.TextTestRunner().run(suite)
