"""ModalGpuEzEmbeddings 유닛 테스트."""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

from langchain_modal_gpu_ez import ModalGpuEzEmbeddings


class TestModalGpuEzEmbeddings:
    """ModalGpuEzEmbeddings 클래스 테스트."""

    def test_default_params(self) -> None:
        emb = ModalGpuEzEmbeddings()
        assert emb.model_id == "BAAI/bge-small-en-v1.5"
        assert emb.gpu == "Local"
        assert emb.model_kwargs == {}

    def test_custom_params(self) -> None:
        emb = ModalGpuEzEmbeddings(
            model_id="sentence-transformers/all-MiniLM-L6-v2",
            gpu="T4",
            model_kwargs={"normalize_embeddings": True},
        )
        assert emb.model_id == "sentence-transformers/all-MiniLM-L6-v2"
        assert emb.gpu == "T4"

    def test_extra_fields_forbidden(self) -> None:
        with pytest.raises(Exception):
            ModalGpuEzEmbeddings(unknown_field="value")  # type: ignore[call-arg]

    @patch("modal_gpu_ez.use")
    def test_embed_documents(self, mock_use: MagicMock) -> None:
        mock_use.return_value = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]

        emb = ModalGpuEzEmbeddings()
        result = emb.embed_documents(["hello", "world"])

        assert len(result) == 2
        assert result[0] == [0.1, 0.2, 0.3]
        assert result[1] == [0.4, 0.5, 0.6]
        mock_use.assert_called_once_with(
            "BAAI/bge-small-en-v1.5",
            "Local",
            input=["hello", "world"],
        )

    @patch("modal_gpu_ez.use")
    def test_embed_query(self, mock_use: MagicMock) -> None:
        mock_use.return_value = [[0.1, 0.2, 0.3]]

        emb = ModalGpuEzEmbeddings()
        result = emb.embed_query("hello")

        assert result == [0.1, 0.2, 0.3]

    @pytest.mark.asyncio
    @patch("modal_gpu_ez.async_use")
    async def test_aembed_documents(self, mock_async_use: MagicMock) -> None:
        mock_async_use.return_value = [[0.1, 0.2], [0.3, 0.4]]

        emb = ModalGpuEzEmbeddings()
        result = await emb.aembed_documents(["hello", "world"])

        assert len(result) == 2
        assert result[0] == [0.1, 0.2]

    @pytest.mark.asyncio
    @patch("modal_gpu_ez.async_use")
    async def test_aembed_query(self, mock_async_use: MagicMock) -> None:
        mock_async_use.return_value = [[0.7, 0.8, 0.9]]

        emb = ModalGpuEzEmbeddings()
        result = await emb.aembed_query("test")

        assert result == [0.7, 0.8, 0.9]
