Ai/integration_tests/base.py

from __future__ import annotations

from datetime import date, timedelta
from typing import Any
import uuid

from django.test import TransactionTestCase
from rest_framework.test import APIClient

from location_data.models import NdviObservation, SoilLocation
from weather.models import WeatherForecast


UNSET = object()


def square_boundary(lat: float, lon: float, delta: float = 0.01) -> dict[str, Any]:
    return {
        "type": "Polygon",
        "coordinates": [
            [
                [lon - delta, lat - delta],
                [lon + delta, lat - delta],
                [lon + delta, lat + delta],
                [lon - delta, lat + delta],
                [lon - delta, lat - delta],
            ]
        ],
    }


class IntegrationAPITestCase(TransactionTestCase):
    reset_sequences = True
    databases = {"default"}

    primary_lat = 35.700000
    primary_lon = 51.400000
    forecast_start = date.today()

    def setUp(self) -> None:
        super().setUp()
        self.client = APIClient()
        self.primary_boundary = square_boundary(self.primary_lat, self.primary_lon)
        self.primary_location = self.create_complete_location(
            lat=self.primary_lat,
            lon=self.primary_lon,
            boundary=self.primary_boundary,
        )
        self.seed_weather_forecasts(self.primary_location, start=self.forecast_start, days=7)
        self.seed_ndvi_observation(self.primary_location)

    def create_complete_location(
        self,
        *,
        lat: float,
        lon: float,
        boundary: dict[str, Any] | None = None,
    ) -> SoilLocation:
        location = SoilLocation.objects.create(
            latitude=f"{lat:.6f}",
            longitude=f"{lon:.6f}",
            farm_boundary=boundary or square_boundary(lat, lon),
        )
        return location

    def seed_weather_forecasts(
        self,
        location: SoilLocation,
        *,
        start: date,
        days: int,
        temperature_base: float = 22.0,
        et0_base: float = 3.4,
    ) -> list[WeatherForecast]:
        forecasts: list[WeatherForecast] = []
        for day_index in range(days):
            forecasts.append(
                WeatherForecast.objects.create(
                    location=location,
                    forecast_date=start + timedelta(days=day_index),
                    temperature_min=12.0 + day_index,
                    temperature_max=temperature_base + day_index,
                    temperature_mean=17.0 + day_index,
                    precipitation=1.2 if day_index % 3 == 0 else 0.0,
                    precipitation_probability=35.0 + day_index,
                    humidity_mean=48.0 + day_index,
                    wind_speed_max=10.0 + day_index,
                    et0=et0_base + (day_index * 0.2),
                    weather_code=0 if day_index == 0 else 2,
                )
            )
        return forecasts

    def seed_ndvi_observation(
        self,
        location: SoilLocation,
        *,
        observation_date: date | None = None,
        mean_ndvi: float = 0.73,
    ) -> NdviObservation:
        return NdviObservation.objects.create(
            location=location,
            observation_date=observation_date or self.forecast_start,
            mean_ndvi=mean_ndvi,
            ndvi_map={"type": "FeatureCollection", "features": []},
            vegetation_health_class="Healthy",
            satellite_source="sentinel-2",
            metadata={"suite": "integration"},
        )

    def create_irrigation_method_via_api(self, name: str, **overrides: Any) -> dict[str, Any]:
        payload = {
            "name": name,
            "category": "localized",
            "description": "Primary drip line for the farm",
            "water_efficiency_percent": 91.0,
            "water_pressure_required": "1.5 bar",
            "flow_rate": "4 l/h",
            "coverage_area": "row-based",
            "soil_type": "loam",
            "climate_suitability": "dry",
        }
        payload.update(overrides)
        response = self.client.post("/api/irrigation/", data=payload, format="json")
        self.assertEqual(response.status_code, 201, response.json())
        return response.json()["data"]

    def create_plant_via_api(self, name: str, **overrides: Any) -> dict[str, Any]:
        payload = {
            "name": name,
            "light": "full sun",
            "watering": "every 2 days",
            "soil": "loamy",
            "temperature": "20-28C",
            "growth_stage": "vegetative",
            "planting_season": "spring",
            "harvest_time": "90 days",
            "spacing": "50 cm",
            "fertilizer": "balanced NPK",
        }
        payload.update(overrides)
        response = self.client.post("/api/plants/", data=payload, format="json")
        self.assertEqual(response.status_code, 201, response.json())
        return response.json()["data"]

    def create_sensor_parameter_via_api(self, **overrides: Any) -> dict[str, Any]:
        payload = {
            "sensor_key": "sensor-7-1",
            "code": "soil_moisture",
            "name_fa": "soil moisture",
            "unit": "%",
            "data_type": "float",
            "metadata": {"min": 0, "max": 100},
        }
        payload.update(overrides)
        response = self.client.post("/api/farm-data/parameters/", data=payload, format="json")
        self.assertEqual(response.status_code, 201, response.json())
        return response.json()["data"]

    def upsert_farm_via_api(
        self,
        *,
        farm_uuid: uuid.UUID,
        plant_ids: list[int] | None = None,
        irrigation_method_id: int | None | object = UNSET,
        sensor_payload: dict[str, Any] | None = None,
        boundary: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "farm_uuid": str(farm_uuid),
            "farm_boundary": boundary or self.primary_boundary,
        }
        if plant_ids is not None:
            payload["plant_ids"] = plant_ids
        if irrigation_method_id is not UNSET:
            payload["irrigation_method_id"] = irrigation_method_id
        if sensor_payload is not None:
            payload["sensor_payload"] = sensor_payload
        response = self.client.post("/api/farm-data/", data=payload, format="json")
        self.assertIn(response.status_code, {200, 201}, response.json())
        return response.json()["data"]