tests/test_postgres_handler.py

import pytest
from unittest.mock import patch, MagicMock
from connectors.postgres_handler import list_postgres_tables, PostgresHandler

@patch('connectors.postgres_handler.create_engine')
@patch('connectors.postgres_handler.inspect')
def test_list_postgres_tables_uses_inspector(mock_inspect, mock_create_engine):
    # Setup our mocks
    mock_engine = MagicMock()
    mock_create_engine.return_value = mock_engine
    
    mock_inspector = MagicMock()
    mock_inspector.get_table_names.return_value = ['users', 'orders']
    mock_inspect.return_value = mock_inspector

    # Execute
    tables = list_postgres_tables('localhost', '5432', 'mydb', 'user', 'pass')

    # Assert
    # We expect the function to return the mocked table names
    assert tables == ['users', 'orders']
    # We expect create_engine to have been called
    mock_create_engine.assert_called_once()
    # Most importantly, we expect it to have used SQLAlchemy inspect
    mock_inspector.get_table_names.assert_called_once_with(schema='public')

@patch('connectors.postgres_handler.create_engine')
@patch('connectors.postgres_handler.inspect')
def test_get_schema_metadata_uses_inspector(mock_inspect, mock_create_engine):
    mock_engine = MagicMock()
    mock_create_engine.return_value = mock_engine
    
    # Mock the SQLAlchemy inspector output
    mock_inspector = MagicMock()
    mock_inspector.get_columns.return_value = [
        {"name": "id", "type": "INTEGER", "comment": None},
        {"name": "email", "type": "VARCHAR", "comment": "user email"}
    ]
    mock_inspect.return_value = mock_inspector

    handler = PostgresHandler()
    
    # Execute with a specific table
    columns = handler.get_schema_metadata('localhost', '5432', 'mydb', 'user', 'pass', 'public.users')

    # Assert expected legacy JSON format
    assert len(columns) == 2
    assert columns[0] == {
        "table_name": "users", 
        "column_name": "id", 
        "data_type": "INTEGER", 
        "column_comment": ""
    }
    assert columns[1] == {
        "table_name": "users", 
        "column_name": "email", 
        "data_type": "VARCHAR", 
        "column_comment": "user email"
    }
    
    mock_inspect.assert_called_once_with(mock_engine)
    mock_inspector.get_columns.assert_called_once_with("users", schema="public")

@patch('connectors.postgres_handler.create_engine')
@patch('connectors.postgres_handler.pd.read_sql')
def test_fetch_data_respects_limit(mock_read_sql, mock_create_engine):
    mock_engine = MagicMock()
    mock_create_engine.return_value = mock_engine
    
    # Return a dummy dataframe
    import pandas as pd
    mock_read_sql.return_value = pd.DataFrame([{"id": 1}])
    
    handler = PostgresHandler()
    
    # Execute with a specific limit
    handler.fetch_data('localhost', '5432', 'mydb', 'user', 'pass', 'users', limit=5)
    
    mock_create_engine.assert_called_once()
    mock_read_sql.assert_called_once()
    
    # Assert that the SQL query passed to read_sql ends with LIMIT 5
    query_passed = mock_read_sql.call_args[0][0]
    assert query_passed.endswith("LIMIT 5")
    assert 'SELECT * FROM "users" LIMIT 5' == query_passed