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