tests/test_api_database_scans.py

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"