test_analytics.py

from __future__ import annotations

from django.contrib.auth.models import User
from django.core.management import call_command
import pytest

from uvdat.core.models import Chart, Dataset, Network, TaskResult
from uvdat.core.tasks.analytics import (
    FloodNetworkFailure,
    FloodSimulation,
    NetworkRecovery,
    analysis_types,
)


@pytest.mark.django_db
def test_rest_list_analysis_types(user, authenticated_api_client, project):
    # TODO: remove this when analytics are no longer hidden from non-superusers
    user.is_superuser = True
    user.save()

    analysis_type_instances = [at() for at in analysis_types if at.is_enabled()]
    resp = authenticated_api_client.get(f"/api/v1/analytics/project/{project.id}/types/")
    data = resp.json()

    assert len(data) == len(analysis_type_instances)
    assert {type_info.get("name") for type_info in data} == {
        i.name for i in analysis_type_instances
    }
    for type_info in data:
        instance = next(iter(i for i in analysis_type_instances if i.name == type_info.get("name")))
        assert type_info.get("description") == instance.description
        assert type_info.get("attribution") == instance.attribution
        assert type_info.get("input_types") == instance.input_types
        assert type_info.get("output_types") == instance.output_types
        input_options = type_info.get("input_options")
        for k, v in instance.get_input_options().items():
            assert len(input_options.get(k)) == len(v)


@pytest.mark.parametrize(
    "task",
    [
        "flood_simulation",
        "flood_network_failure",
        "network_recovery",
        "create_road_network",
    ],
)
@pytest.mark.django_db
def test_rest_run_analysis_task_no_inputs(authenticated_api_client, user, project, task):
    project.set_followers([user])
    resp = authenticated_api_client.post(
        f"/api/v1/analytics/project/{project.id}/types/{task}/run/"
    )
    assert resp.status_code == 400
    assert "not provided" in resp.json()


@pytest.mark.django_db
def test_rest_subscribe_to_running_task(authenticated_api_client, user, project, mailoutbox):
    project.set_followers([user])
    task_result = TaskResult.objects.create(name="Test Task", project=project)
    resp = authenticated_api_client.post(f"/api/v1/analytics/{task_result.id}/subscribe/")
    assert resp.status_code == 200
    task_result.complete()
    message = next(filter(lambda m: m.subject == "GeoDatalytics Task Completed", mailoutbox), None)
    assert message is not None
    assert task_result.name in message.body


@pytest.mark.django_db
def test_rest_subscribe_to_completed_task(authenticated_api_client, user, project):
    project.set_followers([user])
    task_result = TaskResult.objects.create(name="Test Task", project=project)
    task_result.complete()
    resp = authenticated_api_client.post(f"/api/v1/analytics/{task_result.id}/subscribe/")
    assert resp.status_code == 410
    assert resp.json() == "Task already completed. Subscription not applied."


@pytest.mark.slow
@pytest.mark.django_db
def test_flood_analysis_chain(project):
    # ensure a superuser exists
    User.objects.create_superuser("testsuper")

    # ingest necessary objects
    call_command(
        "ingest",
        "./tests/analytics.json",
    )
    network = Network.objects.get(name="MBTA Rapid Transit Network 1")
    chart = Chart.objects.get(name="Charles River Hydrograph (Short Flow Length)")

    # run flood simulation task
    result_1 = FloodSimulation().run_task(
        project=project,
        hydrograph=chart.id,
        time_period="2031-2050",
        annual_probability=0.25,
        potential_evapotranspiration_percentile=50,
        soil_moisture_percentile=50,
        ground_water_percentile=50,
        initial_conditions_id="001",
    )
    result_1.refresh_from_db()
    assert result_1.completed is not None
    assert result_1.error == ""
    assert result_1.outputs is not None

    flood_dataset_id = result_1.outputs.get("flood")
    assert flood_dataset_id is not None
    flood_dataset = Dataset.objects.get(id=flood_dataset_id)
    assert flood_dataset.name == (
        "2031-2050 0.25 Flood Simulation with Charles River "
        "Hydrograph (Short Flow Length), initial condition set 001, and percentiles 50, 50, 50"
    )
    assert flood_dataset.description == "Generated by Flood Simulation Analytics Task"
    assert flood_dataset.category == "flood"

    metadata = flood_dataset.metadata
    assert metadata.get("inputs", {}) == {
        "time_period": "2031-2050",
        "hydrograph": [
            0.006,
            0.026,
            0.066,
            0.111,
            0.138,
            0.143,
            0.128,
            0.102,
            0.08,
            0.054,
            0.038,
            0.03,
            0.021,
            0.016,
            0.012,
            0.008,
            0.006,
            0.005,
            0.003,
            0.002,
            0.002,
            0.001,
            0.001,
            0.001,
        ],
        "pet_percentile": 50,
        "sm_percentile": 50,
        "gw_percentile": 50,
        "annual_probability": 0.25,
        "initial_conditions_id": "001",
    }

    layer = flood_dataset.layers.first()
    assert layer.frames.count() == 24

    # run flood network failure task
    result_2 = FloodNetworkFailure().run_task(
        project=project,
        network=network.id,
        flood_simulation=result_1.id,
        depth_tolerance_meters=0.1,
        station_radius_meters=30,
    )
    result_2.refresh_from_db()
    assert result_2.completed is not None
    assert result_2.error == ""
    assert result_2.outputs is not None

    failures = result_2.outputs.get("failures")
    assert len(failures) == 24

    # run network recovery task
    result_3 = NetworkRecovery().run_task(
        project=project,
        network_failure=result_2.id,
        recovery_mode="degree",
    )
    result_3.refresh_from_db()
    assert result_3.completed is not None
    assert result_3.error == ""
    assert result_3.outputs is not None

    recoveries = result_3.outputs.get("recoveries")
    assert len(recoveries) == 57