Spaces:
Running
Running
File size: 2,607 Bytes
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 | 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"
|