329 lines
11 KiB
Python
329 lines
11 KiB
Python
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from unittest.mock import AsyncMock
|
|
from httpx import ASGITransport, AsyncClient
|
|
|
|
from app.api.v1.datasource_config import get_ai_provider_client
|
|
from app.core.websocket import broadcaster as broadcaster_module
|
|
from app.core.security import get_current_user
|
|
from app.core.target_schema_registry import get_target_schema, list_target_schemas
|
|
from app.main import app
|
|
from app.models.user import User
|
|
from app.models.datasource_config import DataSourceConfig
|
|
from app.services import custom_datasource_runtime
|
|
from app.services.custom_datasource_runtime import run_mapped_websocket_config
|
|
from app.services.datasource_mapping import execute_mapping, persist_mapped_records, redact_for_llm
|
|
|
|
|
|
SAMPLE_AIS = {
|
|
"data": [
|
|
{
|
|
"mmsi": "257123000",
|
|
"latitude": "59.91",
|
|
"longitude": "10.75",
|
|
"speedOverGround": "12.4",
|
|
"timestamp": "2026-04-28T00:00:00Z",
|
|
"api_token": "secret-value",
|
|
}
|
|
]
|
|
}
|
|
|
|
|
|
def test_registry_exposes_v1_target_schemas():
|
|
keys = {schema["key"] for schema in list_target_schemas()}
|
|
|
|
assert {"vessel_ais", "geo_points", "generic_records"}.issubset(keys)
|
|
assert get_target_schema("vessel_ais").destination == "vessel_position"
|
|
|
|
|
|
def test_mapping_engine_maps_and_validates_vessel_ais():
|
|
mapping = {
|
|
"source": {"items_path": "$.data[*]"},
|
|
"fields": {
|
|
"mmsi": {"path": "$.mmsi", "type": "integer"},
|
|
"lat": {"path": "$.latitude", "type": "float"},
|
|
"lon": {"path": "$.longitude", "type": "float"},
|
|
"sog": {"path": "$.speedOverGround", "type": "float"},
|
|
"received_at": {"path": "$.timestamp", "type": "datetime"},
|
|
},
|
|
}
|
|
|
|
result = execute_mapping(SAMPLE_AIS, mapping, "vessel_ais")
|
|
|
|
assert result["mapped_count"] == 1
|
|
assert result["failed_count"] == 0
|
|
assert result["records"][0]["mmsi"] == 257123000
|
|
assert result["records"][0]["lat"] == 59.91
|
|
|
|
|
|
def test_mapping_engine_reports_schema_errors():
|
|
mapping = {
|
|
"source": {"items_path": "$.data[*]"},
|
|
"fields": {
|
|
"mmsi": {"path": "$.mmsi", "type": "integer"},
|
|
"lat": {"path": "$.missing_lat", "type": "float"},
|
|
"lon": {"path": "$.longitude", "type": "float"},
|
|
},
|
|
}
|
|
|
|
result = execute_mapping(SAMPLE_AIS, mapping, "vessel_ais")
|
|
|
|
assert result["mapped_count"] == 0
|
|
assert result["failed_count"] == 1
|
|
assert any("lat" in error for error in result["errors"][0]["errors"])
|
|
|
|
|
|
def test_redact_for_llm_masks_secret_like_fields():
|
|
redacted = redact_for_llm(SAMPLE_AIS)
|
|
|
|
assert redacted["data"][0]["api_token"] == "[REDACTED]"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_mapped_records_writes_generic_records():
|
|
class FakeDB:
|
|
def __init__(self):
|
|
self.added = []
|
|
self.committed = False
|
|
|
|
def add(self, value):
|
|
self.added.append(value)
|
|
|
|
async def commit(self):
|
|
self.committed = True
|
|
|
|
db = FakeDB()
|
|
|
|
count = await persist_mapped_records(
|
|
db,
|
|
datasource_name="custom_weather",
|
|
datasource_config_id=42,
|
|
target_schema="generic_records",
|
|
records=[{"source_id": "row-1", "data": {"temp": 25}}],
|
|
mapping_version=3,
|
|
)
|
|
|
|
assert count == 1
|
|
assert db.committed is True
|
|
assert db.added[0].source == "custom_weather"
|
|
assert db.added[0].data_type == "generic_records"
|
|
assert db.added[0].extra_data["mapping_version"] == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_mapped_vessel_records_writes_raw_and_broadcasts(monkeypatch):
|
|
record_observation = AsyncMock(return_value=object())
|
|
update_health = AsyncMock()
|
|
broadcast_custom = AsyncMock()
|
|
monkeypatch.setattr(
|
|
"app.services.vessel_ais_aggregation.record_vessel_ais_observation",
|
|
record_observation,
|
|
)
|
|
monkeypatch.setattr(
|
|
"app.services.vessel_ais_aggregation.update_ais_source_health",
|
|
update_health,
|
|
)
|
|
monkeypatch.setattr(broadcaster_module, "broadcast_custom", broadcast_custom)
|
|
|
|
class FakeDB:
|
|
def __init__(self):
|
|
self.committed = False
|
|
|
|
async def commit(self):
|
|
self.committed = True
|
|
|
|
db = FakeDB()
|
|
|
|
count = await persist_mapped_records(
|
|
db,
|
|
datasource_name="mock_ais_ws",
|
|
datasource_config_id=42,
|
|
target_schema="vessel_ais",
|
|
records=[
|
|
{
|
|
"mmsi": 999000001,
|
|
"lat": 31.2,
|
|
"lon": 121.4,
|
|
"name": "MOCK VESSEL 001",
|
|
"received_at": "2026-05-01T00:00:00Z",
|
|
}
|
|
],
|
|
mapping_version=1,
|
|
delivery_mode="realtime_stream",
|
|
transport="websocket",
|
|
)
|
|
|
|
assert count == 1
|
|
assert db.committed is True
|
|
record_observation.assert_awaited_once()
|
|
assert record_observation.await_args.kwargs["source"] == "mock_ais_ws"
|
|
assert record_observation.await_args.kwargs["delivery_mode"] == "realtime_stream"
|
|
assert record_observation.await_args.kwargs["transport"] == "websocket"
|
|
update_health.assert_awaited_once()
|
|
broadcast_custom.assert_awaited_once()
|
|
assert broadcast_custom.await_args.args[0] == "vessels"
|
|
assert broadcast_custom.await_args.args[1]["vessels"][0]["mmsi_display"] == "999000001"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_websocket_runner_maps_and_persists_vessel_records(monkeypatch):
|
|
mapping = SimpleNamespace(
|
|
id=7,
|
|
version=2,
|
|
target_schema="vessel_ais",
|
|
mapping_json={
|
|
"source": {"items_path": "$"},
|
|
"fields": {
|
|
"mmsi": {"path": "$.mmsi", "type": "integer"},
|
|
"lat": {"path": "$.lat", "type": "float"},
|
|
"lon": {"path": "$.lon", "type": "float"},
|
|
"name": {"path": "$.name", "type": "string"},
|
|
"received_at": {"path": "$.received_at", "type": "datetime"},
|
|
},
|
|
},
|
|
)
|
|
|
|
class FakeResult:
|
|
def scalar_one_or_none(self):
|
|
return mapping
|
|
|
|
class FakeDB:
|
|
async def execute(self, _stmt):
|
|
return FakeResult()
|
|
|
|
class FakeWebSocket:
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, *_args):
|
|
return None
|
|
|
|
async def send(self, _message):
|
|
return None
|
|
|
|
async def recv(self):
|
|
return (
|
|
'{"type":"vessel","data":{"mmsi":"999000001","name":"MOCK VESSEL 001",'
|
|
'"lat":31.2,"lon":121.4,"received_at":"2026-05-01T00:00:00Z"}}'
|
|
)
|
|
|
|
persist = AsyncMock(return_value=1)
|
|
monkeypatch.setattr(custom_datasource_runtime, "_connect_websocket", AsyncMock(return_value=FakeWebSocket()))
|
|
monkeypatch.setattr(custom_datasource_runtime, "persist_mapped_records", persist)
|
|
|
|
result = await run_mapped_websocket_config(
|
|
FakeDB(),
|
|
DataSourceConfig(
|
|
id=42,
|
|
name="mock_ais_ws",
|
|
source_type="websocket",
|
|
endpoint="ws://localhost:8787/ais",
|
|
auth_type="none",
|
|
headers={},
|
|
config={"ws_message_path": "$.data", "debug_max_messages": 1},
|
|
),
|
|
)
|
|
|
|
assert result["status"] == "success"
|
|
assert result["messages_seen"] == 1
|
|
assert result["written_count"] == 1
|
|
persist.assert_awaited_once()
|
|
assert persist.await_args.kwargs["datasource_name"] == "mock_ais_ws"
|
|
assert persist.await_args.kwargs["records"][0]["mmsi"] == 999000001
|
|
assert persist.await_args.kwargs["delivery_mode"] == "realtime_stream"
|
|
assert persist.await_args.kwargs["transport"] == "websocket"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mapping_preview_api_uses_deterministic_engine():
|
|
def override_get_current_user():
|
|
return User(
|
|
id=1,
|
|
username="testuser",
|
|
email="test@example.com",
|
|
password_hash="hashed",
|
|
role="admin",
|
|
is_active=True,
|
|
)
|
|
|
|
app.dependency_overrides = {get_current_user: override_get_current_user}
|
|
transport = ASGITransport(app=app)
|
|
try:
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
response = await client.post(
|
|
"/api/v1/datasources/mappings/preview",
|
|
json={
|
|
"sample_payload": SAMPLE_AIS,
|
|
"target_schema": "vessel_ais",
|
|
"mapping_json": {
|
|
"source": {"items_path": "$.data[*]"},
|
|
"fields": {
|
|
"mmsi": {"path": "$.mmsi", "type": "integer"},
|
|
"lat": {"path": "$.latitude", "type": "float"},
|
|
"lon": {"path": "$.longitude", "type": "float"},
|
|
},
|
|
},
|
|
},
|
|
)
|
|
finally:
|
|
app.dependency_overrides.clear()
|
|
|
|
assert response.status_code == 200
|
|
payload = response.json()
|
|
assert payload["success"] is True
|
|
assert payload["preview"]["records"][0]["mmsi"] == 257123000
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mapping_propose_api_redacts_sample_before_ai():
|
|
seen_context = {}
|
|
|
|
class FakeAIClient:
|
|
async def analyze(self, request, request_id=None):
|
|
seen_context.update(request.context)
|
|
return SimpleNamespace(
|
|
content=(
|
|
'{"source":{"items_path":"$.data[*]"},"fields":{'
|
|
'"mmsi":{"path":"$.mmsi","type":"integer"},'
|
|
'"lat":{"path":"$.latitude","type":"float"},'
|
|
'"lon":{"path":"$.longitude","type":"float"}}}'
|
|
)
|
|
)
|
|
|
|
def override_get_current_user():
|
|
return User(
|
|
id=1,
|
|
username="testuser",
|
|
email="test@example.com",
|
|
password_hash="hashed",
|
|
role="admin",
|
|
is_active=True,
|
|
)
|
|
|
|
def override_ai_client():
|
|
return FakeAIClient()
|
|
|
|
app.dependency_overrides = {
|
|
get_current_user: override_get_current_user,
|
|
get_ai_provider_client: override_ai_client,
|
|
}
|
|
transport = ASGITransport(app=app)
|
|
try:
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
response = await client.post(
|
|
"/api/v1/datasources/mappings/propose",
|
|
json={
|
|
"sample_payload": SAMPLE_AIS,
|
|
"target_schema": "vessel_ais",
|
|
"use_ai": True,
|
|
},
|
|
)
|
|
finally:
|
|
app.dependency_overrides.clear()
|
|
|
|
assert response.status_code == 200
|
|
payload = response.json()
|
|
assert payload["mapping_json"]["meta"]["generated_by"] == "ai_provider"
|
|
assert seen_context["sample_payload"]["data"][0]["api_token"] == "[REDACTED]"
|