ktx/python/klo-daemon/tests/test_code_execution.py

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any

import numpy as np
import orjson
import pandas as pd
import pytest

from klo_daemon.code_execution import (
    ExecuteCodeRequest,
    create_scratchpad_helpers,
    detect_visualizations,
    dumps_numpy_json,
    execute_code_response,
)


@dataclass
class FakeResponse:
    json_payload: dict[str, Any] | None = None
    content: bytes = b""
    headers: dict[str, str] | None = None

    def raise_for_status(self) -> None:
        return None

    def json(self) -> dict[str, Any]:
        return self.json_payload or {}


class FakeHttpClient:
    def __init__(self) -> None:
        self.posts: list[dict[str, Any]] = []
        self.gets: list[dict[str, Any]] = []

    def post(
        self,
        url: str,
        data: bytes,
        headers: dict[str, str],
        timeout: int,
    ) -> FakeResponse:
        self.posts.append(
            {
                "url": url,
                "data": orjson.loads(data),
                "headers": headers,
                "timeout": timeout,
            }
        )
        return FakeResponse(json_payload={"filename": "saved.json"})

    def get(
        self,
        url: str,
        headers: dict[str, str],
        timeout: int,
    ) -> FakeResponse:
        self.gets.append({"url": url, "headers": headers, "timeout": timeout})
        return FakeResponse(
            content=b"value,name\n1.25,alpha\n",
            headers={"content-type": "text/csv; charset=utf-8"},
        )


def test_execute_code_response_captures_console_result_and_strips_ansi() -> None:
    response = execute_code_response(
        ExecuteCodeRequest(
            code='print("\\x1b[31mhello\\x1b[0m")\nresult = {"value": 3}',
        ),
        nest_api_url=None,
        auth_header=None,
    )

    assert response.result == {"value": 3}
    assert response.console_output == "\x1b[31mhello\x1b[0m\n"
    assert "=== Console Output ===" in response.formatted_result
    assert "hello" in response.formatted_result
    assert "\x1b" not in response.formatted_result
    assert "=== Result ===" in response.formatted_result


def test_execute_code_response_returns_message_when_result_is_absent() -> None:
    response = execute_code_response(
        ExecuteCodeRequest(code='print("ran")'),
        nest_api_url=None,
        auth_header=None,
    )

    assert response.result is None
    assert (
        response.message == "Code executed successfully but no result variable was set"
    )
    assert response.console_output == "ran\n"
    assert "=== Message ===" in response.formatted_result


def test_execute_code_response_detects_visualization_records() -> None:
    response = execute_code_response(
        ExecuteCodeRequest(
            code="result = "
            + json.dumps(
                {
                    "type": "visualization",
                    "vis_type": "bar",
                    "config": {"title": "Revenue"},
                    "data": [{"month": "Jan", "revenue": 10}],
                    "title": "Revenue",
                }
            ),
        ),
        nest_api_url=None,
        auth_header=None,
    )

    assert response.visualizations is not None
    assert len(response.visualizations) == 1
    assert response.visualizations[0].vis_type == "bar"
    assert response.visualizations[0].title == "Revenue"


def test_detect_visualizations_filters_mixed_lists() -> None:
    visualizations = detect_visualizations(
        [
            {"type": "note", "text": "skip"},
            {
                "type": "visualization",
                "vis_type": "table",
                "config": {"title": "Rows"},
                "data": [{"row": 1}],
            },
        ]
    )

    assert visualizations == [
        {
            "type": "visualization",
            "vis_type": "table",
            "config": {"title": "Rows"},
            "data": [{"row": 1}],
        }
    ]


def test_scratchpad_and_visualization_helpers_serialize_numpy_scalars() -> None:
    client = FakeHttpClient()
    save_df, read_file, save_viz = create_scratchpad_helpers(
        nest_api_url="http://nest",
        auth_header="Bearer token",
        source_id="source_123",
        message_id="message_456",
        http_client=client,
    )

    df = pd.DataFrame({"value": [np.float64(1.25)]})
    assert save_df(df, filename="df.json") == "1 rows saved to saved.json"

    read_df = read_file("input.csv")
    assert read_df.to_dict(orient="records") == [{"value": 1.25, "name": "alpha"}]

    viz_ref = save_viz(
        vis_type="bar",
        config={"title": "Test", "x": "a", "y": np.float64(2.5)},
        data=[{"a": "row1", "b": np.float64(3.75)}],
    )
    assert viz_ref == "![viz](saved.json)"

    assert (
        client.posts[0]["url"] == "http://nest/private_api/scratchpad/source_123/files"
    )
    assert client.posts[0]["data"]["data"][0]["value"] == 1.25
    assert (
        client.gets[0]["url"]
        == "http://nest/private_api/scratchpad/source_123/files/input.csv?format=raw"
    )
    assert client.posts[1]["url"] == "http://nest/private_api/visualizations/source_123"
    assert client.posts[1]["data"]["config"]["y"] == 2.5
    assert client.posts[1]["data"]["data"][0]["b"] == 3.75


def test_scratchpad_helpers_require_app_context_only_when_called() -> None:
    save_df, read_file, save_viz = create_scratchpad_helpers(
        nest_api_url=None,
        auth_header=None,
        source_id=None,
        message_id=None,
    )

    with pytest.raises(ValueError, match="required for scratchpad operations"):
        save_df(pd.DataFrame({"value": [1]}), filename="df.json")

    with pytest.raises(ValueError, match="required for scratchpad operations"):
        read_file("df.csv")

    with pytest.raises(ValueError, match="required for visualization operations"):
        save_viz("bar", {"title": "Chart"}, [{"value": 1}])


def test_dumps_numpy_json_serializes_numpy_values() -> None:
    rendered = dumps_numpy_json(
        {
            "scalar": np.float64(1.5),
            "array": np.array([1, 2, 3]),
        }
    )

    assert orjson.loads(rendered) == {"scalar": 1.5, "array": [1, 2, 3]}