import pytest import os os.environ["SUPABASE_URL"] = "http://localhost:8000" os.environ["SUPABASE_ANON_KEY"] = "dummy-key" os.environ["OPENAI_API_KEY"] = "dummy-key" from fastapi.testclient import TestClient from unittest.mock import patch, MagicMock from api import app import pandas as pd client = TestClient(app) @patch('api.classifier.get_postgres_data') @patch('api.classifier.pg_handler.get_schema_metadata') @patch('classifier_manager.metadata_scanner.MetadataScanner.scan_schema') @patch('api._uid_from_auth') @patch('db.supabase_client.persist_scan_results') def test_connect_postgresql_hybrid_mode( mock_persist, mock_uid, mock_scan_metadata, mock_get_schema, mock_get_data ): mock_uid.return_value = "test-user-123" # Mock data sample returned from fetch_data (5 rows limit simulation) mock_get_data.return_value = pd.DataFrame([{"id": 1, "name": "John Doe"}]) # Mock schema metadata fetched from handler mock_get_schema.return_value = [ {"column_name": "name", "data_type": "VARCHAR"} ] # Mock MetadataScanner response mock_scan_metadata.return_value = { "metadata": { "flagged_columns": [ {"column_name": "name", "matched_rule": "PERSON"} ] } } payload = { "host": "localhost", "port": "5432", "database": "mydb", "user": "postgres", "password": "password", "table": "users", "scan_mode": "metadata_and_sampling" } response = client.post("/api/connect/postgresql", json=payload, headers={"Authorization": "Bearer token123"}) assert response.status_code == 200 res_json = response.json() # 1. Verify limit=5 was requested mock_get_data.assert_called_once_with( "localhost", "5432", "mydb", "postgres", "password", "users", limit=5 ) # 2. Verify schema metadata was fetched mock_get_schema.assert_called_once_with( "localhost", "5432", "mydb", "postgres", "password", "users" ) # 3. Verify MetadataScanner was invoked mock_scan_metadata.assert_called_once() # 4. Verify hybrid JSON response structure assert "data" in res_json assert len(res_json["data"]) == 1 assert "metadata" in res_json assert "flagged_columns" in res_json["metadata"] assert res_json["metadata"]["flagged_columns"][0]["column_name"] == "name" # 5. Verify persist logic received the correct mode mock_persist.assert_called_once() assert mock_persist.call_args[1]["scan_mode"] == "metadata_and_sampling"