Spaces:
Running
Running
File size: 3,157 Bytes
af03eab 468ae07 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | 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
|