#
# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
# or more contributor license agreements. Licensed under the Elastic License 2.0;
# you may not use this file except in compliance with the Elastic License 2.0.
#
"""Tests the PostgreSQL database source class methods"""

import ssl
from contextlib import asynccontextmanager
from datetime import datetime
from decimal import Decimal
from ipaddress import (
    IPv4Address,
    IPv4Interface,
    IPv4Network,
    IPv6Address,
    IPv6Interface,
    IPv6Network,
)
from unittest.mock import ANY, Mock, patch
from uuid import UUID

import pytest
from asyncpg.types import (
    BitString,
    Box,
    Circle,
    Line,
    LineSegment,
    Path,
    Point,
    Polygon,
)
from connectors_sdk.filtering.validation import Filter, SyncRuleValidationResult
from connectors_sdk.utils import iso_utc
from freezegun import freeze_time
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.ext.asyncio.engine import AsyncEngine

from connectors.sources.postgresql import (
    PostgreSQLAdvancedRulesValidator,
    PostgreSQLClient,
    PostgreSQLDataSource,
    PostgreSQLQueries,
)
from tests.sources.support import create_source

ADVANCED_SNIPPET = "advanced_snippet"
POSTGRESQL_CONNECTION_STRING = (
    "postgresql+asyncpg://admin:changme@127.0.0.1:5432/testdb"
)
SCHEMA = "public"
TABLE = "emp_table"
CUSTOMER_TABLE = "customer"

TIME = iso_utc()

ID_ONE = "id1"
ID_TWO = "id2"
ID_THREE = "id3"


@asynccontextmanager
async def create_postgresql_source():
    async with create_source(
        PostgreSQLDataSource,
        host="127.0.0.1",
        port="9090",
        database="xe",
        tables="*",
        schema=SCHEMA,
    ) as source:
        yield source


class MockSsl:
    """This class contains methods which returns dummy ssl context"""

    def load_verify_locations(self, cadata):
        """This method verify locations"""
        pass


class ConnectionAsync:
    """This class creates dummy connection with database and return dummy cursor"""

    async def __aenter__(self):
        """Make a dummy database connection and return it"""
        return self

    async def __aexit__(self, exception_type, exception_value, exception_traceback):
        """Make sure the dummy database connection gets closed"""
        pass

    async def execute(self, query):
        """This method returns dummy cursor"""
        return CursorAsync(query=query)


class CursorAsync:
    """This class contains methods which returns dummy response"""

    async def __aenter__(self):
        """Make a dummy database connection and return it"""
        return self

    def __init__(self, *args, **kw):
        """Setup dummy cursor"""
        self.query = kw["query"]
        self.first_call = True

    def keys(self):
        """Return Columns of table

        Returns:
            list: List of columns
        """
        return ["ids", "names"]

    def fetchmany(self, size):
        """This method returns response of fetchmany

        Args:
            size (int): Number of rows

        Returns:
            list: List of rows
        """
        if self.first_call:
            self.first_call = False

            self.query = str(self.query)
            query_object = PostgreSQLQueries()
            if self.query == query_object.all_tables(database="xe", schema=SCHEMA):
                return [(TABLE,)]
            elif self.query == query_object.table_data_count(
                schema=SCHEMA, table=TABLE
            ):
                return [(10,)]
            elif self.query == query_object.table_primary_key(
                schema=SCHEMA, table=TABLE
            ):
                return [("ids",)]
            elif self.query == query_object.table_primary_key(
                schema=SCHEMA, table=CUSTOMER_TABLE
            ):
                return [("ids",)]
            elif self.query == query_object.table_last_update_time(
                schema=SCHEMA, table=TABLE
            ):
                return [("2023-02-21T08:37:15+00:00",)]
            elif self.query == query_object.table_last_update_time(
                schema=SCHEMA, table=CUSTOMER_TABLE
            ):
                return [("2023-02-21T08:37:15+00:00",)]
            elif self.query.lower() == "select * from customer":
                return [(1, "customer_1"), (2, "customer_2")]
            elif self.query == query_object.ping():
                return [(2,)]
            else:
                return [
                    (
                        1,
                        "abcd",
                    ),
                    (
                        2,
                        "xyz",
                    ),
                ]
        return []

    async def __aexit__(self, exception_type, exception_value, exception_traceback):
        """Make sure the dummy database connection gets closed"""
        pass


def test_get_connect_args():
    """This function test _get_connect_args with dummy certificate"""
    # Setup
    client = PostgreSQLClient(
        host="",
        port="",
        user="",
        password="",
        database="",
        schema="",
        tables="*",
        ssl_enabled=True,
        ssl_ca="-----BEGIN CERTIFICATE----- Certificate -----END CERTIFICATE-----",
        logger_=None,
    )

    # Execute
    with patch.object(ssl, "create_default_context", return_value=MockSsl()):
        client._get_connect_args()


@pytest.mark.asyncio
async def test_postgresql_ping():
    # Setup
    async with create_postgresql_source() as source:
        with patch.object(AsyncEngine, "connect", return_value=ConnectionAsync()):
            await source.ping()

        await source.close()


@pytest.mark.asyncio
async def test_ping():
    async with create_postgresql_source() as source:
        with patch.object(AsyncEngine, "connect", return_value=ConnectionAsync()):
            await source.ping()


@pytest.mark.asyncio
@patch("connectors.utils.time_to_sleep_between_retries", Mock(return_value=0))
async def test_ping_negative():
    with pytest.raises(Exception):
        async with create_source(PostgreSQLDataSource, port=5432) as source:
            with patch.object(AsyncEngine, "connect", side_effect=Exception()):
                await source.ping()


@pytest.mark.parametrize(
    "advanced_rules, expected_validation_result",
    [
        (
            # valid: empty array should be valid
            [],
            SyncRuleValidationResult.valid_result(
                SyncRuleValidationResult.ADVANCED_RULES
            ),
        ),
        (
            # valid: empty object should also be valid -> default value in Kibana
            {},
            SyncRuleValidationResult.valid_result(
                SyncRuleValidationResult.ADVANCED_RULES
            ),
        ),
        (
            # valid: valid queries
            [
                {
                    "tables": ["emp_table"],
                    "query": "select * from emp_table",
                }
            ],
            SyncRuleValidationResult.valid_result(
                SyncRuleValidationResult.ADVANCED_RULES
            ),
        ),
        (
            # invalid: tables not present in database
            [
                {
                    "tables": ["table_name"],
                    "query": "select * from table_name",
                }
            ],
            SyncRuleValidationResult(
                SyncRuleValidationResult.ADVANCED_RULES,
                is_valid=False,
                validation_message=ANY,
            ),
        ),
        (
            # invalid: tables key missing
            [{"query": "select * from table_name"}],
            SyncRuleValidationResult(
                SyncRuleValidationResult.ADVANCED_RULES,
                is_valid=False,
                validation_message=ANY,
            ),
        ),
        (
            # invalid: invalid key
            [
                {
                    "tables": "table_name",
                    "query": "select * from table_name",
                }
            ],
            SyncRuleValidationResult(
                SyncRuleValidationResult.ADVANCED_RULES,
                is_valid=False,
                validation_message=ANY,
            ),
        ),
        (
            # invalid: tables can be empty
            [
                {
                    "tables": [],
                    "query": "select * from table_name",
                }
            ],
            SyncRuleValidationResult(
                SyncRuleValidationResult.ADVANCED_RULES,
                is_valid=False,
                validation_message=ANY,
            ),
        ),
    ],
)
@pytest.mark.asyncio
async def test_advanced_rules_validation(advanced_rules, expected_validation_result):
    async with create_source(
        PostgreSQLDataSource, database="xe", tables="*", schema="public", port=5432
    ) as source:
        with patch.object(AsyncEngine, "connect", return_value=ConnectionAsync()):
            validation_result = await PostgreSQLAdvancedRulesValidator(source).validate(
                advanced_rules
            )

            assert validation_result == expected_validation_result


@pytest.mark.parametrize(
    "advanced_rules, id_in_source, expected_validation_result",
    [
        (
            # valid: empty array should be valid
            [],
            [ID_ONE],
            SyncRuleValidationResult.valid_result(
                SyncRuleValidationResult.ADVANCED_RULES
            ),
        ),
        (
            # valid: empty object should also be valid -> default value in Kibana
            {},
            [ID_ONE],
            SyncRuleValidationResult.valid_result(
                SyncRuleValidationResult.ADVANCED_RULES
            ),
        ),
        (
            # valid: valid queries
            [
                {
                    "tables": ["emp_table"],
                    "query": "select * from emp_table",
                    "id_columns": [ID_ONE],
                }
            ],
            [ID_ONE],
            SyncRuleValidationResult.valid_result(
                SyncRuleValidationResult.ADVANCED_RULES
            ),
        ),
        (
            # valid: id_columns field missing
            [
                {
                    "tables": ["emp_table"],
                    "query": "select * from emp_table",
                    "id_columns": [ID_ONE],
                }
            ],
            [ID_ONE],
            SyncRuleValidationResult.valid_result(
                SyncRuleValidationResult.ADVANCED_RULES
            ),
        ),
        (
            # invalid: tables not present in database
            [
                {
                    "tables": ["table_name"],
                    "query": "select * from table_name",
                    "id_columns": [ID_ONE],
                }
            ],
            [ID_ONE],
            SyncRuleValidationResult(
                SyncRuleValidationResult.ADVANCED_RULES,
                is_valid=False,
                validation_message=ANY,
            ),
        ),
        (
            # invalid: tables key missing
            [{"query": "select * from table_name", "id_columns": [ID_ONE]}],
            [ID_ONE],
            SyncRuleValidationResult(
                SyncRuleValidationResult.ADVANCED_RULES,
                is_valid=False,
                validation_message=ANY,
            ),
        ),
        (
            # invalid: invalid key
            [
                {
                    "tables": "table_name",
                    "query": "select * from table_name",
                    "id_columns": [ID_ONE],
                }
            ],
            [],
            SyncRuleValidationResult(
                SyncRuleValidationResult.ADVANCED_RULES,
                is_valid=False,
                validation_message=ANY,
            ),
        ),
        (
            # invalid: tables can be empty
            [
                {
                    "tables": [],
                    "query": "select * from table_name",
                    "id_columns": [ID_ONE],
                }
            ],
            [ID_ONE],
            SyncRuleValidationResult(
                SyncRuleValidationResult.ADVANCED_RULES,
                is_valid=False,
                validation_message=ANY,
            ),
        ),
    ],
)
@pytest.mark.asyncio
async def test_advanced_rules_validation_when_id_in_source_available(
    advanced_rules, id_in_source, expected_validation_result
):
    async with create_source(
        PostgreSQLDataSource, database="xe", tables="*", schema="public", port=5432
    ) as source:
        with patch.object(AsyncEngine, "connect", return_value=ConnectionAsync()):
            validation_result = await PostgreSQLAdvancedRulesValidator(source).validate(
                advanced_rules
            )

            assert validation_result == expected_validation_result


@freeze_time(TIME)
@pytest.mark.asyncio
async def test_get_docs():
    # Setup
    async with create_postgresql_source() as source:
        with patch.object(AsyncEngine, "connect", return_value=ConnectionAsync()):
            source.engine = create_async_engine(POSTGRESQL_CONNECTION_STRING)
            actual_response = []
            expected_response = [
                {
                    "public_emp_table_ids": 1,
                    "public_emp_table_names": "abcd",
                    "_id": "xe_public_emp_table_1",
                    "_timestamp": "2023-02-21T08:37:15+00:00",
                    "database": "xe",
                    "table": "emp_table",
                    "schema": "public",
                },
                {
                    "public_emp_table_ids": 2,
                    "public_emp_table_names": "xyz",
                    "_id": "xe_public_emp_table_2",
                    "_timestamp": "2023-02-21T08:37:15+00:00",
                    "database": "xe",
                    "table": "emp_table",
                    "schema": "public",
                },
            ]

            # Execute
            async for doc in source.get_docs():
                actual_response.append(doc[0])

            # Assert
            assert actual_response == expected_response


@freeze_time(TIME)
@pytest.mark.parametrize(
    "filtering, expected_response",
    [
        # Configured valid query
        (
            Filter(
                {
                    ADVANCED_SNIPPET: {
                        "value": [
                            {
                                "tables": ["emp_table"],
                                "query": "select * from emp_table",
                            },
                        ]
                    }
                }
            ),
            [
                {
                    "public_emp_table_ids": 1,
                    "public_emp_table_names": "abcd",
                    "_id": "xe_public_emp_table_1",
                    "_timestamp": "2023-02-21T08:37:15+00:00",
                    "database": "xe",
                    "table": ["emp_table"],
                    "schema": "public",
                },
                {
                    "public_emp_table_ids": 2,
                    "public_emp_table_names": "xyz",
                    "_id": "xe_public_emp_table_2",
                    "_timestamp": "2023-02-21T08:37:15+00:00",
                    "database": "xe",
                    "table": ["emp_table"],
                    "schema": "public",
                },
            ],
        ),
        (
            # Configured multiple rules
            Filter(
                {
                    ADVANCED_SNIPPET: {
                        "value": [
                            {
                                "tables": ["emp_table"],
                                "query": "select * from emp_table",
                            },
                            {"tables": ["customer"], "query": "select * from customer"},
                        ]
                    }
                }
            ),
            [
                {
                    "public_emp_table_ids": 1,
                    "public_emp_table_names": "abcd",
                    "_id": "xe_public_emp_table_1",
                    "_timestamp": "2023-02-21T08:37:15+00:00",
                    "database": "xe",
                    "table": ["emp_table"],
                    "schema": "public",
                },
                {
                    "public_emp_table_ids": 2,
                    "public_emp_table_names": "xyz",
                    "_id": "xe_public_emp_table_2",
                    "_timestamp": "2023-02-21T08:37:15+00:00",
                    "database": "xe",
                    "table": ["emp_table"],
                    "schema": "public",
                },
                {
                    "public_customer_ids": 1,
                    "public_customer_names": "customer_1",
                    "_id": "xe_public_customer_1",
                    "_timestamp": "2023-02-21T08:37:15+00:00",
                    "database": "xe",
                    "table": ["customer"],
                    "schema": "public",
                },
                {
                    "public_customer_ids": 2,
                    "public_customer_names": "customer_2",
                    "_id": "xe_public_customer_2",
                    "_timestamp": "2023-02-21T08:37:15+00:00",
                    "database": "xe",
                    "table": ["customer"],
                    "schema": "public",
                },
            ],
        ),
        (
            # Configured with id_columns using mixed case - tests case sensitivity
            Filter(
                {
                    ADVANCED_SNIPPET: {
                        "value": [
                            {
                                "tables": ["emp_table"],
                                "query": "select * from emp_table",
                                "id_columns": [
                                    "IDS"
                                ],  # uppercase to test case insensitivity
                            },
                        ]
                    }
                }
            ),
            [
                {
                    "public_emp_table_ids": 1,
                    "public_emp_table_names": "abcd",
                    "_id": "xe_public_emp_table_1",
                    "_timestamp": "2023-02-21T08:37:15+00:00",
                    "database": "xe",
                    "table": ["emp_table"],
                    "schema": "public",
                },
                {
                    "public_emp_table_ids": 2,
                    "public_emp_table_names": "xyz",
                    "_id": "xe_public_emp_table_2",
                    "_timestamp": "2023-02-21T08:37:15+00:00",
                    "database": "xe",
                    "table": ["emp_table"],
                    "schema": "public",
                },
            ],
        ),
    ],
)
@pytest.mark.asyncio
async def test_get_docs_with_advanced_rules(filtering, expected_response):
    async with create_source(
        PostgreSQLDataSource,
        database="xe",
        tables="*",
        schema="public",
        port=5432,
    ) as source:
        with patch.object(AsyncEngine, "connect", return_value=ConnectionAsync()):
            actual_response = []

            async for doc in source.get_docs(filtering=filtering):
                actual_response.append(doc[0])

            assert actual_response == expected_response


def test_table_primary_key_query():
    """Test that the table_primary_key method generates the correct SQL query"""
    # Setup
    queries = PostgreSQLQueries()
    schema = "test_schema"
    table = "test_table"

    # Execute
    query = queries.table_primary_key(schema=schema, table=table)

    # Assert
    expected_query = (
        "SELECT a.attname AS c "
        "FROM pg_index i "
        "JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) "
        "JOIN pg_class t ON t.oid = i.indrelid "
        "JOIN pg_constraint c ON c.conindid = i.indexrelid "
        'WHERE i.indrelid = \'"test_schema"."test_table"\'::regclass '
        "AND t.relkind = 'r' "
        "AND c.contype = 'p' "
        "ORDER BY array_position(i.indkey, a.attnum)"
    )
    assert query == expected_query


def test_table_primary_key_query_with_special_characters():
    """Test that the table_primary_key method handles special characters in schema and table names"""
    # Setup
    queries = PostgreSQLQueries()
    schema = "test-schema"
    table = 'test_table-with.special"chars'

    # Execute
    query = queries.table_primary_key(schema=schema, table=table)

    # Assert
    expected_query = (
        "SELECT a.attname AS c "
        "FROM pg_index i "
        "JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey) "
        "JOIN pg_class t ON t.oid = i.indrelid "
        "JOIN pg_constraint c ON c.conindid = i.indexrelid "
        'WHERE i.indrelid = \'"test-schema"."test_table-with.special"chars"\'::regclass '
        "AND t.relkind = 'r' "
        "AND c.contype = 'p' "
        "ORDER BY array_position(i.indkey, a.attnum)"
    )
    assert query == expected_query


@pytest.mark.asyncio
async def test_get_table_primary_key():
    """Test that the get_table_primary_key method correctly processes query results"""
    # Setup
    async with create_postgresql_source() as source:
        with patch.object(AsyncEngine, "connect", return_value=ConnectionAsync()):
            # Execute
            primary_keys = await source.postgresql_client.get_table_primary_key(TABLE)

            # Assert
            assert primary_keys == ["ids"]


@pytest.mark.asyncio
async def test_serialize_ipv4_address():
    """Test that IPv4Address objects are correctly serialized to strings"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "name": "test",
            "ip_address": IPv4Address("192.168.1.1"),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["ip_address"] == "192.168.1.1"


@pytest.mark.asyncio
async def test_serialize_ipv6_address():
    """Test that IPv6Address objects are correctly serialized to strings"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "name": "test",
            "ip_address": IPv6Address("2001:db8::1"),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["ip_address"] == "2001:db8::1"


@pytest.mark.asyncio
async def test_serialize_ip_address_in_list():
    """Test that IP addresses nested in lists are correctly serialized"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "ip_addresses": [
                IPv4Address("10.0.0.1"),
                IPv4Address("10.0.0.2"),
                IPv6Address("fe80::1"),
            ],
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["ip_addresses"] == ["10.0.0.1", "10.0.0.2", "fe80::1"]
        assert all(isinstance(ip, str) for ip in serialized["ip_addresses"])


@pytest.mark.asyncio
async def test_serialize_ip_address_in_nested_dict():
    """Test that IP addresses nested in dictionaries are correctly serialized"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "connection": {
                "source_ip": IPv4Address("192.168.1.100"),
                "dest_ip": IPv4Address("8.8.8.8"),
            },
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["connection"]["source_ip"] == "192.168.1.100"
        assert serialized["connection"]["dest_ip"] == "8.8.8.8"


@pytest.mark.asyncio
async def test_serialize_mixed_types_with_ip():
    """Test that serialization handles mixed data types including IP addresses"""
    # Setup
    async with create_postgresql_source() as source:
        timestamp = datetime(2023, 1, 15, 10, 30, 45)
        doc = {
            "id": 1,
            "ip_address": IPv4Address("10.30.0.9"),
            "timestamp": timestamp,
            "price": Decimal("99.99"),
            "data": bytes("test", "utf-8"),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["ip_address"] == "10.30.0.9"
        assert serialized["timestamp"] == timestamp.isoformat()
        assert serialized["price"] == 99.99
        assert serialized["data"] == "test"


@pytest.mark.asyncio
async def test_serialize_ipv4_network():
    """Test that IPv4Network objects (CIDR type) are correctly serialized to strings"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "network": IPv4Network("192.168.0.0/24"),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["network"] == "192.168.0.0/24"


@pytest.mark.asyncio
async def test_serialize_ipv6_network():
    """Test that IPv6Network objects (CIDR type) are correctly serialized to strings"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "network": IPv6Network("2001:db8::/32"),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["network"] == "2001:db8::/32"


@pytest.mark.asyncio
async def test_serialize_ipv4_interface():
    """Test that IPv4Interface objects (INET type) are correctly serialized to strings"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "interface": IPv4Interface("192.168.1.1/24"),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["interface"] == "192.168.1.1/24"


@pytest.mark.asyncio
async def test_serialize_ipv6_interface():
    """Test that IPv6Interface objects (INET type) are correctly serialized to strings"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "interface": IPv6Interface("2001:db8::1/64"),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["interface"] == "2001:db8::1/64"


@pytest.mark.asyncio
async def test_serialize_all_network_types():
    """Test that all PostgreSQL network types are correctly serialized"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "ipv4_addr": IPv4Address("192.168.1.1"),
            "ipv6_addr": IPv6Address("2001:db8::1"),
            "ipv4_iface": IPv4Interface("10.0.0.1/8"),
            "ipv6_iface": IPv6Interface("fe80::1/64"),
            "ipv4_net": IPv4Network("172.16.0.0/12"),
            "ipv6_net": IPv6Network("2001:db8::/32"),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["ipv4_addr"] == "192.168.1.1"
        assert serialized["ipv6_addr"] == "2001:db8::1"
        assert serialized["ipv4_iface"] == "10.0.0.1/8"
        assert serialized["ipv6_iface"] == "fe80::1/64"
        assert serialized["ipv4_net"] == "172.16.0.0/12"
        assert serialized["ipv6_net"] == "2001:db8::/32"
        assert all(isinstance(v, str) for k, v in serialized.items() if k != "id")


@pytest.mark.asyncio
async def test_serialize_uuid():
    """Test that UUID objects are correctly serialized to strings"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "uuid_col": UUID("a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11"),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["uuid_col"] == "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11"


@pytest.mark.asyncio
async def test_serialize_uuid_in_list():
    """Test that UUIDs nested in lists are correctly serialized"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "uuid_array": [
                UUID("550e8400-e29b-41d4-a716-446655440000"),
                UUID("f47ac10b-58cc-4372-a567-0e02b2c3d479"),
            ],
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["uuid_array"] == [
            "550e8400-e29b-41d4-a716-446655440000",
            "f47ac10b-58cc-4372-a567-0e02b2c3d479",
        ]
        assert all(isinstance(uuid, str) for uuid in serialized["uuid_array"])


@pytest.mark.asyncio
async def test_serialize_mixed_special_types():
    """Test serialization of multiple special types together"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "ip_addr": IPv4Address("192.168.1.1"),
            "ip_network": IPv4Network("10.0.0.0/8"),
            "uuid_col": UUID("a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11"),
            "mixed_array": [
                IPv4Address("10.0.0.1"),
                UUID("550e8400-e29b-41d4-a716-446655440000"),
            ],
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["ip_addr"] == "192.168.1.1"
        assert serialized["ip_network"] == "10.0.0.0/8"
        assert serialized["uuid_col"] == "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11"
        assert serialized["mixed_array"] == [
            "10.0.0.1",
            "550e8400-e29b-41d4-a716-446655440000",
        ]


@pytest.mark.asyncio
async def test_serialize_geometric_point():
    """Test serialization of PostgreSQL POINT type"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "location": Point(1.5, 2.5),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["location"] == "(1.5, 2.5)"


@pytest.mark.asyncio
async def test_serialize_geometric_line():
    """Test serialization of PostgreSQL LINE type"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "boundary": Line(1, -1, 0),  # x - y = 0
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["boundary"] == "(1, -1, 0)"


@pytest.mark.asyncio
async def test_serialize_geometric_lseg():
    """Test serialization of PostgreSQL LSEG (line segment) type"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "segment": LineSegment((0, 0), (1, 1)),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["segment"] == "[(0.0, 0.0), (1.0, 1.0)]"


@pytest.mark.asyncio
async def test_serialize_geometric_box():
    """Test serialization of PostgreSQL BOX type"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "area": Box((2, 2), (0, 0)),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["area"] == "[(2.0, 2.0), (0.0, 0.0)]"


@pytest.mark.asyncio
async def test_serialize_geometric_path():
    """Test serialization of PostgreSQL PATH type"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "route": Path((0, 0), (1, 1), (2, 0)),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["route"] == "open [(0.0, 0.0), (1.0, 1.0), (2.0, 0.0)]"


@pytest.mark.asyncio
async def test_serialize_geometric_polygon():
    """Test serialization of PostgreSQL POLYGON type"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "shape": Polygon((0, 0), (1, 0), (1, 1), (0, 1)),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["shape"] == "[(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 1.0)]"


@pytest.mark.asyncio
async def test_serialize_geometric_circle():
    """Test serialization of PostgreSQL CIRCLE type"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "zone": Circle((0, 0), 5),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["zone"] == "((0, 0), 5)"


@pytest.mark.asyncio
async def test_serialize_bitstring():
    """Test serialization of PostgreSQL BIT/VARBIT types"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            "flags": BitString("10101010"),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert
        assert serialized["flags"] == "1010 1010"


@pytest.mark.asyncio
async def test_serialize_all_special_types():
    """Test serialization of all PostgreSQL special types together"""
    # Setup
    async with create_postgresql_source() as source:
        doc = {
            "id": 1,
            # Network types
            "ip_inet": IPv4Address("192.168.1.1"),
            "ip_cidr": IPv4Network("10.0.0.0/8"),
            # UUID
            "uuid_col": UUID("a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11"),
            # Geometric types
            "point_col": Point(1.5, 2.5),
            "line_col": Line(1, -1, 0),
            "lseg_col": LineSegment((0, 0), (1, 1)),
            "box_col": Box((2, 2), (0, 0)),
            "path_col": Path((0, 0), (1, 1), (2, 0)),
            "polygon_col": Polygon((0, 0), (1, 0), (1, 1), (0, 1)),
            "circle_col": Circle((0, 0), 5),
            # BitString
            "bit_col": BitString("10101010"),
        }

        # Execute
        serialized = source.serialize(doc)

        # Assert - check exact serialized values
        assert serialized["ip_inet"] == "192.168.1.1"
        assert serialized["ip_cidr"] == "10.0.0.0/8"
        assert serialized["uuid_col"] == "a0eebc99-9c0b-4ef8-bb6d-6bb9bd380a11"
        assert serialized["point_col"] == "(1.5, 2.5)"
        assert serialized["line_col"] == "(1, -1, 0)"
        assert serialized["lseg_col"] == "[(0.0, 0.0), (1.0, 1.0)]"
        assert serialized["box_col"] == "[(2.0, 2.0), (0.0, 0.0)]"
        assert serialized["path_col"] == "open [(0.0, 0.0), (1.0, 1.0), (2.0, 0.0)]"
        assert (
            serialized["polygon_col"]
            == "[(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 1.0)]"
        )
        assert serialized["circle_col"] == "((0, 0), 5)"
        assert serialized["bit_col"] == "1010 1010"
