Presa-Mano/tests/test_sam21_client.py

"""Unit tests for Sam21Client."""
from __future__ import annotations

from unittest.mock import patch

import httpx
import numpy as np
import pytest

from presa_mano.clients.sam21 import Sam21Client, Sam21Error, Sam21Unavailable


def _intrinsics():
    return {"fx": 600.0, "fy": 600.0, "cx": 320.0, "cy": 240.0,
            "width": 640, "height": 480}


def _ok_response(points_mm):
    body = {
        "success": True,
        "data": {
            "n_points": len(points_mm),
            "n_pixels_in_mask": len(points_mm),
            "n_pixels_with_valid_depth": len(points_mm),
            "elapsed_ms": {"segment": 100.0, "project": 5.0, "total": 105.0},
            "points": [list(p) + [0.0, 0.0, 0.0] for p in points_mm],
        },
        "error": None,
    }
    return httpx.Response(200, json=body)


def _err_response(status: int, code: str, message: str):
    body = {"success": False, "data": None,
            "error": {"code": code, "message": message}}
    return httpx.Response(status, json=body)


def test_segment_pointcloud_happy_path():
    points_mm = [(100.0, 200.0, 1000.0), (110.0, 200.0, 1010.0)]
    response = _ok_response(points_mm)
    client = Sam21Client("http://sam21:8002", timeout_s=5.0)
    with patch("httpx.Client.post", return_value=response) as mock_post:
        out = client.segment_pointcloud(
            rgb_png=b"FAKE_PNG", depth_png=b"FAKE_DEPTH",
            intrinsics=_intrinsics(), bbox=(10, 20, 100, 200),
        )
    assert mock_post.call_count == 1
    posted_url = mock_post.call_args.args[0]
    assert posted_url.endswith("/api/segment_pointcloud")
    body = mock_post.call_args.kwargs["json"]
    assert body["bbox"] == [10, 20, 100, 200]
    assert body["depth_scale_mm"] == 1.0
    assert body["return_format"] == "json"
    assert body["return_mask"] is False
    assert isinstance(body["rgb_png_base64"], str) and len(body["rgb_png_base64"]) > 0
    assert isinstance(body["depth_png_base64"], str) and len(body["depth_png_base64"]) > 0
    assert body["intrinsics"] == {
        "fx": 600.0, "fy": 600.0, "cx": 320.0, "cy": 240.0,
        "width": 640, "height": 480,
    }
    # Output is in meters: mm/1000.
    assert out.shape == (2, 3)
    assert out.dtype == np.float64
    assert np.allclose(out, np.array([[0.1, 0.2, 1.0], [0.11, 0.2, 1.01]]))


def test_segment_pointcloud_connect_error_raises_unavailable():
    client = Sam21Client("http://sam21:8002")
    with patch("httpx.Client.post", side_effect=httpx.ConnectError("refused")):
        with pytest.raises(Sam21Unavailable):
            client.segment_pointcloud(
                rgb_png=b"x", depth_png=b"y",
                intrinsics=_intrinsics(), bbox=(0, 0, 10, 10),
            )


def test_segment_pointcloud_timeout_raises_unavailable():
    client = Sam21Client("http://sam21:8002")
    with patch("httpx.Client.post", side_effect=httpx.TimeoutException("timeout")):
        with pytest.raises(Sam21Unavailable):
            client.segment_pointcloud(
                rgb_png=b"x", depth_png=b"y",
                intrinsics=_intrinsics(), bbox=(0, 0, 10, 10),
            )


def test_segment_pointcloud_model_not_ready_raises_unavailable():
    """503 MODEL_NOT_READY is transient → treated as unavailable (fallback)."""
    response = _err_response(503, "MODEL_NOT_READY", "loading")
    client = Sam21Client("http://sam21:8002")
    with patch("httpx.Client.post", return_value=response):
        with pytest.raises(Sam21Unavailable):
            client.segment_pointcloud(
                rgb_png=b"x", depth_png=b"y",
                intrinsics=_intrinsics(), bbox=(0, 0, 10, 10),
            )


def test_segment_pointcloud_segmentation_empty_raises_error():
    """422 SEGMENTATION_EMPTY is a semantic outcome, not transport."""
    response = _err_response(422, "SEGMENTATION_EMPTY", "no mask")
    client = Sam21Client("http://sam21:8002")
    with patch("httpx.Client.post", return_value=response):
        with pytest.raises(Sam21Error) as info:
            client.segment_pointcloud(
                rgb_png=b"x", depth_png=b"y",
                intrinsics=_intrinsics(), bbox=(0, 0, 10, 10),
            )
    assert "SEGMENTATION_EMPTY" in str(info.value)


def test_segment_pointcloud_no_valid_depth_raises_error():
    response = _err_response(422, "NO_VALID_DEPTH", "mask found but no depth")
    client = Sam21Client("http://sam21:8002")
    with patch("httpx.Client.post", return_value=response):
        with pytest.raises(Sam21Error) as info:
            client.segment_pointcloud(
                rgb_png=b"x", depth_png=b"y",
                intrinsics=_intrinsics(), bbox=(0, 0, 10, 10),
            )
    assert "NO_VALID_DEPTH" in str(info.value)


def test_segment_pointcloud_empty_points_returns_empty_array():
    """Service returned success but 0 points — return (0, 3) array."""
    response = _ok_response([])
    client = Sam21Client("http://sam21:8002")
    with patch("httpx.Client.post", return_value=response):
        out = client.segment_pointcloud(
            rgb_png=b"x", depth_png=b"y",
            intrinsics=_intrinsics(), bbox=(0, 0, 10, 10),
        )
    assert out.shape == (0, 3)