"""Tests for the GrillyOptimum package.

Covers:
  - VulkanConfig: defaults, serialization, deserialization, round-trip
  - VulkanModelForCausalLM: class interface, generate method logic
  - Pipeline: create_text_generation_pipeline, VulkanTextGenerationPipeline

All tests are CPU-only and mock heavy dependencies (grillyinference,
transformers) so the suite runs without GPU hardware or model weights.
"""

from __future__ import annotations

import importlib
import sys
from dataclasses import fields
from unittest.mock import MagicMock, patch, PropertyMock

import numpy as np
import pytest

# ---------------------------------------------------------------------------
# VulkanConfig tests
# ---------------------------------------------------------------------------
from grillyoptimum.configuration import VulkanConfig


class TestVulkanConfig:
    """Tests for VulkanConfig dataclass."""

    # -- default values -----------------------------------------------------

    def test_default_dtype(self):
        cfg = VulkanConfig()
        assert cfg.dtype == "fp16"

    def test_default_use_vulkan(self):
        cfg = VulkanConfig()
        assert cfg.use_vulkan is True

    def test_default_page_size(self):
        cfg = VulkanConfig()
        assert cfg.page_size == 256

    def test_default_raw_window(self):
        cfg = VulkanConfig()
        assert cfg.raw_window == 2048

    def test_default_enable_h2o(self):
        cfg = VulkanConfig()
        assert cfg.enable_h2o is False

    def test_default_h2o_lambda(self):
        cfg = VulkanConfig()
        assert cfg.h2o_lambda == pytest.approx(0.0002)

    def test_default_enable_vsa(self):
        cfg = VulkanConfig()
        assert cfg.enable_vsa is False

    def test_default_enable_quantization(self):
        cfg = VulkanConfig()
        assert cfg.enable_quantization is False

    def test_default_quantize_group_size(self):
        cfg = VulkanConfig()
        assert cfg.quantize_group_size == 64

    def test_default_max_batch_size(self):
        cfg = VulkanConfig()
        assert cfg.max_batch_size == 1

    def test_default_device(self):
        cfg = VulkanConfig()
        assert cfg.device == "vulkan"

    # -- to_dict() ----------------------------------------------------------

    def test_to_dict_returns_dict(self):
        cfg = VulkanConfig()
        d = cfg.to_dict()
        assert isinstance(d, dict)

    def test_to_dict_contains_all_fields(self):
        cfg = VulkanConfig()
        d = cfg.to_dict()
        field_names = {f.name for f in fields(VulkanConfig)}
        assert set(d.keys()) == field_names

    def test_to_dict_values_match_defaults(self):
        cfg = VulkanConfig()
        d = cfg.to_dict()
        assert d["dtype"] == "fp16"
        assert d["use_vulkan"] is True
        assert d["page_size"] == 256
        assert d["raw_window"] == 2048
        assert d["enable_h2o"] is False
        assert d["h2o_lambda"] == pytest.approx(0.0002)
        assert d["enable_vsa"] is False
        assert d["enable_quantization"] is False
        assert d["quantize_group_size"] == 64
        assert d["max_batch_size"] == 1
        assert d["device"] == "vulkan"

    def test_to_dict_reflects_custom_values(self):
        cfg = VulkanConfig(dtype="fp32", page_size=512, enable_h2o=True)
        d = cfg.to_dict()
        assert d["dtype"] == "fp32"
        assert d["page_size"] == 512
        assert d["enable_h2o"] is True

    # -- from_dict() --------------------------------------------------------

    def test_from_dict_creates_instance(self):
        d = {"dtype": "fp32", "use_vulkan": False, "page_size": 128}
        cfg = VulkanConfig.from_dict(d)
        assert isinstance(cfg, VulkanConfig)
        assert cfg.dtype == "fp32"
        assert cfg.use_vulkan is False
        assert cfg.page_size == 128

    def test_from_dict_ignores_unknown_keys(self):
        d = {"dtype": "fp16", "unknown_key": 999, "extra": "value"}
        cfg = VulkanConfig.from_dict(d)
        assert cfg.dtype == "fp16"
        assert not hasattr(cfg, "unknown_key")
        assert not hasattr(cfg, "extra")

    def test_from_dict_uses_defaults_for_missing_keys(self):
        cfg = VulkanConfig.from_dict({"dtype": "fp32"})
        # dtype is overridden
        assert cfg.dtype == "fp32"
        # everything else falls back to defaults
        assert cfg.use_vulkan is True
        assert cfg.page_size == 256
        assert cfg.raw_window == 2048

    def test_from_dict_empty_dict_gives_defaults(self):
        cfg = VulkanConfig.from_dict({})
        default = VulkanConfig()
        assert cfg.to_dict() == default.to_dict()

    # -- custom config values -----------------------------------------------

    def test_custom_config_all_fields(self):
        cfg = VulkanConfig(
            dtype="fp32",
            use_vulkan=False,
            page_size=512,
            raw_window=4096,
            enable_h2o=True,
            h2o_lambda=0.001,
            enable_vsa=True,
            enable_quantization=True,
            quantize_group_size=128,
            max_batch_size=4,
            device="cpu",
        )
        assert cfg.dtype == "fp32"
        assert cfg.use_vulkan is False
        assert cfg.page_size == 512
        assert cfg.raw_window == 4096
        assert cfg.enable_h2o is True
        assert cfg.h2o_lambda == pytest.approx(0.001)
        assert cfg.enable_vsa is True
        assert cfg.enable_quantization is True
        assert cfg.quantize_group_size == 128
        assert cfg.max_batch_size == 4
        assert cfg.device == "cpu"

    # -- round-trip ---------------------------------------------------------

    def test_round_trip_default(self):
        original = VulkanConfig()
        restored = VulkanConfig.from_dict(original.to_dict())
        assert original.to_dict() == restored.to_dict()

    def test_round_trip_custom(self):
        original = VulkanConfig(
            dtype="fp32",
            use_vulkan=False,
            page_size=1024,
            raw_window=8192,
            enable_h2o=True,
            h2o_lambda=0.05,
            enable_vsa=True,
            enable_quantization=True,
            quantize_group_size=32,
            max_batch_size=8,
            device="cpu",
        )
        restored = VulkanConfig.from_dict(original.to_dict())
        assert original.to_dict() == restored.to_dict()

    def test_round_trip_preserves_types(self):
        original = VulkanConfig(h2o_lambda=0.123)
        d = original.to_dict()
        restored = VulkanConfig.from_dict(d)
        assert type(restored.dtype) is str
        assert type(restored.use_vulkan) is bool
        assert type(restored.page_size) is int
        assert type(restored.h2o_lambda) is float
        assert restored.h2o_lambda == pytest.approx(0.123)


# ---------------------------------------------------------------------------
# VulkanModelForCausalLM tests
# ---------------------------------------------------------------------------
from grillyoptimum.modeling import VulkanModelForCausalLM


class TestVulkanModelForCausalLM:
    """Tests for VulkanModelForCausalLM class interface and generate logic."""

    # -- class structure ----------------------------------------------------

    def test_class_exists(self):
        assert VulkanModelForCausalLM is not None

    def test_has_from_pretrained(self):
        assert hasattr(VulkanModelForCausalLM, "from_pretrained")
        assert callable(VulkanModelForCausalLM.from_pretrained)

    def test_has_generate(self):
        assert hasattr(VulkanModelForCausalLM, "generate")
        assert callable(VulkanModelForCausalLM.generate)

    def test_has_call(self):
        assert hasattr(VulkanModelForCausalLM, "__call__")

    def test_has_device_property(self):
        assert isinstance(
            getattr(VulkanModelForCausalLM, "device", None),
            property,
        )

    def test_has_memory_footprint(self):
        assert hasattr(VulkanModelForCausalLM, "memory_footprint")
        assert callable(VulkanModelForCausalLM.memory_footprint)

    # -- constructor --------------------------------------------------------

    def test_init_stores_model_and_config(self):
        mock_model = MagicMock()
        mock_config = MagicMock()
        m = VulkanModelForCausalLM(mock_model, mock_config)
        assert m._model is mock_model
        assert m.config is mock_config

    def test_init_default_vulkan_config(self):
        m = VulkanModelForCausalLM(MagicMock(), MagicMock())
        assert isinstance(m.vulkan_config, VulkanConfig)

    def test_init_custom_vulkan_config(self):
        vc = VulkanConfig(dtype="fp32", page_size=512)
        m = VulkanModelForCausalLM(MagicMock(), MagicMock(), vulkan_config=vc)
        assert m.vulkan_config is vc
        assert m.vulkan_config.dtype == "fp32"
        assert m.vulkan_config.page_size == 512

    def test_init_kv_cache_none(self):
        m = VulkanModelForCausalLM(MagicMock(), MagicMock())
        assert m._kv_cache is None

    # -- device property ----------------------------------------------------

    def test_device_returns_vulkan_config_device(self):
        vc = VulkanConfig(device="cpu")
        m = VulkanModelForCausalLM(MagicMock(), MagicMock(), vulkan_config=vc)
        assert m.device == "cpu"

    def test_device_default_is_vulkan(self):
        m = VulkanModelForCausalLM(MagicMock(), MagicMock())
        assert m.device == "vulkan"

    # -- generate input validation ------------------------------------------

    def test_generate_raises_on_none_input_ids(self):
        m = VulkanModelForCausalLM(MagicMock(), MagicMock())
        with pytest.raises(ValueError, match="input_ids is required"):
            m.generate(input_ids=None)

    def test_generate_raises_on_batch_size_gt_1(self):
        m = VulkanModelForCausalLM(MagicMock(), MagicMock())
        ids = np.array([[1, 2], [3, 4]], dtype=np.int32)
        with pytest.raises(ValueError, match="batch_size=1"):
            # Must patch grillyinference imports inside generate
            with patch.dict(sys.modules, {
                "grillyinference": MagicMock(),
                "grillyinference.inference": MagicMock(),
                "grillyinference.inference.generate": MagicMock(),
            }):
                m.generate(input_ids=ids)

    def test_generate_converts_list_to_ndarray(self):
        """generate() should accept a plain list and convert it."""
        mock_model = MagicMock()
        mock_config = MagicMock()

        # Set up mock returns
        fake_logits = np.zeros((1, 5, 100), dtype=np.float32)
        mock_model.forward.return_value = fake_logits
        mock_model.decode_step.return_value = fake_logits

        mock_kv_cache_cls = MagicMock()

        mock_sample_fn = MagicMock(return_value=128009)  # stop token
        mock_gen_module = MagicMock()
        mock_gen_module._sample_top_k_top_p = mock_sample_fn
        mock_gen_module.LLAMA3_STOP_TOKENS = {128001, 128008, 128009}

        with patch.dict(sys.modules, {
            "grillyinference": MagicMock(KVCache=mock_kv_cache_cls),
            "grillyinference.inference": MagicMock(),
            "grillyinference.inference.generate": mock_gen_module,
        }):
            m = VulkanModelForCausalLM(mock_model, mock_config)
            result = m.generate(input_ids=[1, 2, 3], max_new_tokens=5)

        assert isinstance(result, np.ndarray)
        assert result.ndim == 2

    # -- generate with mocks ------------------------------------------------

    def test_generate_calls_forward_and_decode_step(self):
        """Verify generate calls _model.forward for prefill and
        _model.decode_step for subsequent tokens."""
        mock_model = MagicMock()
        mock_config = MagicMock()

        vocab_size = 50
        fake_logits = np.random.randn(1, 3, vocab_size).astype(np.float32)
        mock_model.forward.return_value = fake_logits

        step_logits = np.random.randn(1, 1, vocab_size).astype(np.float32)
        mock_model.decode_step.return_value = step_logits

        # The sample function returns token 42, then a stop token
        call_count = {"n": 0}
        stop_token = 128009

        def fake_sample(logits, temperature, top_k, top_p):
            call_count["n"] += 1
            if call_count["n"] >= 3:
                return stop_token
            return 42

        mock_gen_module = MagicMock()
        mock_gen_module._sample_top_k_top_p = fake_sample
        mock_gen_module.LLAMA3_STOP_TOKENS = {128001, 128008, 128009}

        mock_kv_cache_cls = MagicMock()

        with patch.dict(sys.modules, {
            "grillyinference": MagicMock(KVCache=mock_kv_cache_cls),
            "grillyinference.inference": MagicMock(),
            "grillyinference.inference.generate": mock_gen_module,
        }):
            m = VulkanModelForCausalLM(mock_model, mock_config)
            input_ids = np.array([[10, 20, 30]], dtype=np.int32)
            result = m.generate(input_ids=input_ids, max_new_tokens=10)

        # forward called once for prefill
        mock_model.forward.assert_called_once()
        # decode_step called for each generated token after the first sample
        assert mock_model.decode_step.call_count == 2  # tokens 42, 42, then stop

        # result includes original tokens + generated
        assert result.shape[0] == 1
        assert list(result[0][:3]) == [10, 20, 30]
        assert 42 in result[0]

    def test_generate_max_length_overrides_max_new_tokens(self):
        """When max_length is provided it should override max_new_tokens."""
        mock_model = MagicMock()
        mock_config = MagicMock()

        fake_logits = np.zeros((1, 3, 50), dtype=np.float32)
        mock_model.forward.return_value = fake_logits
        mock_model.decode_step.return_value = fake_logits

        # Always return stop token so generation ends immediately
        stop_token = 128009
        mock_gen_module = MagicMock()
        mock_gen_module._sample_top_k_top_p = MagicMock(return_value=stop_token)
        mock_gen_module.LLAMA3_STOP_TOKENS = {128001, 128008, 128009}

        mock_kv_cache_cls = MagicMock()

        with patch.dict(sys.modules, {
            "grillyinference": MagicMock(KVCache=mock_kv_cache_cls),
            "grillyinference.inference": MagicMock(),
            "grillyinference.inference.generate": mock_gen_module,
        }):
            m = VulkanModelForCausalLM(mock_model, mock_config)
            input_ids = np.array([[1, 2, 3]], dtype=np.int32)
            # max_length=5 with input length 3 => max_new_tokens should be 2
            result = m.generate(input_ids=input_ids, max_length=5)

        # Generation should still succeed (stop token hit immediately)
        assert result.shape == (1, 3)

    def test_generate_greedy_sets_temperature_zero(self):
        """When do_sample=False, temperature should be set to 0."""
        mock_model = MagicMock()
        mock_config = MagicMock()

        fake_logits = np.zeros((1, 3, 50), dtype=np.float32)
        mock_model.forward.return_value = fake_logits

        stop_token = 128009
        captured_temps = []

        def capture_sample(logits, temperature, top_k, top_p):
            captured_temps.append(temperature)
            return stop_token

        mock_gen_module = MagicMock()
        mock_gen_module._sample_top_k_top_p = capture_sample
        mock_gen_module.LLAMA3_STOP_TOKENS = {128001, 128008, 128009}

        mock_kv_cache_cls = MagicMock()

        with patch.dict(sys.modules, {
            "grillyinference": MagicMock(KVCache=mock_kv_cache_cls),
            "grillyinference.inference": MagicMock(),
            "grillyinference.inference.generate": mock_gen_module,
        }):
            m = VulkanModelForCausalLM(mock_model, mock_config)
            input_ids = np.array([[1, 2, 3]], dtype=np.int32)
            m.generate(input_ids=input_ids, do_sample=False, max_new_tokens=5)

        assert len(captured_temps) >= 1
        assert all(t == 0.0 for t in captured_temps)

    def test_call_delegates_to_generate(self):
        """__call__ should forward to generate()."""
        mock_model = MagicMock()
        mock_config = MagicMock()
        m = VulkanModelForCausalLM(mock_model, mock_config)

        with patch.object(m, "generate", return_value=np.array([[1]])) as mock_gen:
            ids = np.array([[1, 2]], dtype=np.int32)
            m(ids, max_new_tokens=10)
            mock_gen.assert_called_once_with(ids, max_new_tokens=10)

    def test_memory_footprint_delegates(self):
        mock_model = MagicMock()
        mock_model.memory_footprint.return_value = {"total_mb": 1024}
        m = VulkanModelForCausalLM(mock_model, MagicMock())
        result = m.memory_footprint()
        assert result == {"total_mb": 1024}
        mock_model.memory_footprint.assert_called_once()

    # -- from_pretrained (mocked) -------------------------------------------

    def test_from_pretrained_returns_instance(self):
        """from_pretrained should return a VulkanModelForCausalLM."""
        mock_llama_model = MagicMock()
        mock_llama_config = MagicMock()

        mock_llama_cls = MagicMock()
        mock_llama_cls.from_pretrained.return_value = mock_llama_model

        mock_config_cls = MagicMock()
        mock_config_cls.from_pretrained.return_value = mock_llama_config

        mock_gi = MagicMock()
        mock_gi.LlamaForCausalLM = mock_llama_cls
        mock_gi.LlamaConfig = mock_config_cls

        with patch.dict(sys.modules, {"grillyinference": mock_gi}):
            # Reload modeling so the import inside from_pretrained picks up the mock
            import grillyoptimum.modeling as mod
            result = mod.VulkanModelForCausalLM.from_pretrained("fake/model")

        assert isinstance(result, VulkanModelForCausalLM)
        assert result.config is mock_llama_config
        assert result._model is mock_llama_model
        assert isinstance(result.vulkan_config, VulkanConfig)

    def test_from_pretrained_passes_dtype(self):
        """from_pretrained should pass dtype to VulkanConfig and model load."""
        mock_llama_cls = MagicMock()
        mock_config_cls = MagicMock()

        mock_gi = MagicMock()
        mock_gi.LlamaForCausalLM = mock_llama_cls
        mock_gi.LlamaConfig = mock_config_cls

        with patch.dict(sys.modules, {"grillyinference": mock_gi}):
            import grillyoptimum.modeling as mod
            result = mod.VulkanModelForCausalLM.from_pretrained(
                "fake/model", dtype="fp32",
            )

        assert result.vulkan_config.dtype == "fp32"
        mock_llama_cls.from_pretrained.assert_called_once_with(
            "fake/model", dtype="fp32",
        )

    def test_from_pretrained_uses_explicit_vulkan_config(self):
        """When vulkan_config is passed, it should be used as-is."""
        mock_llama_cls = MagicMock()
        mock_config_cls = MagicMock()

        mock_gi = MagicMock()
        mock_gi.LlamaForCausalLM = mock_llama_cls
        mock_gi.LlamaConfig = mock_config_cls

        vc = VulkanConfig(dtype="fp32", page_size=1024)

        with patch.dict(sys.modules, {"grillyinference": mock_gi}):
            import grillyoptimum.modeling as mod
            result = mod.VulkanModelForCausalLM.from_pretrained(
                "fake/model", vulkan_config=vc,
            )

        assert result.vulkan_config is vc
        assert result.vulkan_config.page_size == 1024


# ---------------------------------------------------------------------------
# Pipeline tests
# ---------------------------------------------------------------------------


class TestPipeline:
    """Tests for pipeline.py: create_text_generation_pipeline and
    VulkanTextGenerationPipeline."""

    # -- function existence -------------------------------------------------

    def test_create_text_generation_pipeline_exists(self):
        from grillyoptimum.pipeline import create_text_generation_pipeline
        assert callable(create_text_generation_pipeline)

    # -- VulkanTextGenerationPipeline structure (via creating one) -----------

    def test_pipeline_returns_callable(self):
        """create_text_generation_pipeline should return a callable object."""
        mock_llama_cls = MagicMock()
        mock_text_gen = MagicMock()
        mock_tokenizer = MagicMock()

        mock_gi = MagicMock()
        mock_gi.LlamaForCausalLM = mock_llama_cls
        mock_gi.TextGenerator = mock_text_gen

        with patch.dict(sys.modules, {
            "grillyinference": mock_gi,
        }):
            from grillyoptimum.pipeline import create_text_generation_pipeline
            pipe = create_text_generation_pipeline(
                "fake/model", tokenizer=mock_tokenizer, dtype="fp16",
            )

        assert callable(pipe)

    def test_pipeline_has_generator_and_tokenizer(self):
        """The returned pipeline should store _generator and _tokenizer."""
        mock_llama_cls = MagicMock()
        mock_text_gen_cls = MagicMock()
        mock_text_gen_instance = MagicMock()
        mock_text_gen_cls.return_value = mock_text_gen_instance
        mock_tokenizer = MagicMock()

        mock_gi = MagicMock()
        mock_gi.LlamaForCausalLM = mock_llama_cls
        mock_gi.TextGenerator = mock_text_gen_cls

        with patch.dict(sys.modules, {
            "grillyinference": mock_gi,
        }):
            from grillyoptimum.pipeline import create_text_generation_pipeline
            pipe = create_text_generation_pipeline(
                "fake/model", tokenizer=mock_tokenizer,
            )

        assert hasattr(pipe, "_generator")
        assert hasattr(pipe, "_tokenizer")
        assert pipe._tokenizer is mock_tokenizer

    def test_pipeline_auto_loads_tokenizer_when_none(self):
        """If no tokenizer is given, AutoTokenizer.from_pretrained is called."""
        mock_llama_cls = MagicMock()
        mock_text_gen_cls = MagicMock()
        mock_auto_tok = MagicMock()
        mock_auto_tok.from_pretrained.return_value = MagicMock()

        mock_gi = MagicMock()
        mock_gi.LlamaForCausalLM = mock_llama_cls
        mock_gi.TextGenerator = mock_text_gen_cls

        mock_transformers = MagicMock()
        mock_transformers.AutoTokenizer = mock_auto_tok

        with patch.dict(sys.modules, {
            "grillyinference": mock_gi,
            "transformers": mock_transformers,
        }):
            from grillyoptimum.pipeline import create_text_generation_pipeline
            pipe = create_text_generation_pipeline("fake/model", tokenizer=None)

        mock_auto_tok.from_pretrained.assert_called_once_with("fake/model")

    def test_pipeline_call_single_string(self):
        """Calling the pipeline with a single string should return a list
        of dicts with 'generated_text' and 'prompt' keys."""
        mock_llama_cls = MagicMock()
        mock_text_gen_cls = MagicMock()
        mock_generator = MagicMock()
        mock_generator.generate.return_value = "Hello world response"
        mock_text_gen_cls.return_value = mock_generator
        mock_tokenizer = MagicMock()

        mock_gi = MagicMock()
        mock_gi.LlamaForCausalLM = mock_llama_cls
        mock_gi.TextGenerator = mock_text_gen_cls

        with patch.dict(sys.modules, {
            "grillyinference": mock_gi,
        }):
            from grillyoptimum.pipeline import create_text_generation_pipeline
            pipe = create_text_generation_pipeline(
                "fake/model", tokenizer=mock_tokenizer,
            )
            results = pipe("Hello world")

        assert isinstance(results, list)
        assert len(results) == 1
        assert results[0]["generated_text"] == "Hello world response"
        assert results[0]["prompt"] == "Hello world"

    def test_pipeline_call_list_of_strings(self):
        """Calling with a list of strings should produce one result per prompt."""
        mock_llama_cls = MagicMock()
        mock_text_gen_cls = MagicMock()
        mock_generator = MagicMock()
        mock_generator.generate.side_effect = ["Response A", "Response B"]
        mock_text_gen_cls.return_value = mock_generator
        mock_tokenizer = MagicMock()

        mock_gi = MagicMock()
        mock_gi.LlamaForCausalLM = mock_llama_cls
        mock_gi.TextGenerator = mock_text_gen_cls

        with patch.dict(sys.modules, {
            "grillyinference": mock_gi,
        }):
            from grillyoptimum.pipeline import create_text_generation_pipeline
            pipe = create_text_generation_pipeline(
                "fake/model", tokenizer=mock_tokenizer,
            )
            results = pipe(["Prompt A", "Prompt B"])

        assert len(results) == 2
        assert results[0]["prompt"] == "Prompt A"
        assert results[0]["generated_text"] == "Response A"
        assert results[1]["prompt"] == "Prompt B"
        assert results[1]["generated_text"] == "Response B"

    def test_pipeline_passes_generation_kwargs(self):
        """Generation kwargs (max_new_tokens, temperature, top_k, top_p)
        should be forwarded to generator.generate()."""
        mock_llama_cls = MagicMock()
        mock_text_gen_cls = MagicMock()
        mock_generator = MagicMock()
        mock_generator.generate.return_value = "output"
        mock_text_gen_cls.return_value = mock_generator
        mock_tokenizer = MagicMock()

        mock_gi = MagicMock()
        mock_gi.LlamaForCausalLM = mock_llama_cls
        mock_gi.TextGenerator = mock_text_gen_cls

        with patch.dict(sys.modules, {
            "grillyinference": mock_gi,
        }):
            from grillyoptimum.pipeline import create_text_generation_pipeline
            pipe = create_text_generation_pipeline(
                "fake/model", tokenizer=mock_tokenizer,
            )
            pipe(
                "test prompt",
                max_new_tokens=256,
                temperature=0.5,
                top_k=40,
                top_p=0.95,
            )

        mock_generator.generate.assert_called_once_with(
            "test prompt",
            max_tokens=256,
            temperature=0.5,
            top_k=40,
            top_p=0.95,
        )

    def test_pipeline_class_name(self):
        """The inner pipeline class should be named VulkanTextGenerationPipeline."""
        mock_llama_cls = MagicMock()
        mock_text_gen_cls = MagicMock()
        mock_tokenizer = MagicMock()

        mock_gi = MagicMock()
        mock_gi.LlamaForCausalLM = mock_llama_cls
        mock_gi.TextGenerator = mock_text_gen_cls

        with patch.dict(sys.modules, {
            "grillyinference": mock_gi,
        }):
            from grillyoptimum.pipeline import create_text_generation_pipeline
            pipe = create_text_generation_pipeline(
                "fake/model", tokenizer=mock_tokenizer,
            )

        assert type(pipe).__name__ == "VulkanTextGenerationPipeline"

    def test_pipeline_loads_model_with_dtype(self):
        """create_text_generation_pipeline should pass dtype to
        LlamaForCausalLM.from_pretrained."""
        mock_llama_cls = MagicMock()
        mock_text_gen_cls = MagicMock()
        mock_tokenizer = MagicMock()

        mock_gi = MagicMock()
        mock_gi.LlamaForCausalLM = mock_llama_cls
        mock_gi.TextGenerator = mock_text_gen_cls

        with patch.dict(sys.modules, {
            "grillyinference": mock_gi,
        }):
            from grillyoptimum.pipeline import create_text_generation_pipeline
            create_text_generation_pipeline(
                "fake/model", tokenizer=mock_tokenizer, dtype="fp32",
            )

        mock_llama_cls.from_pretrained.assert_called_once_with(
            "fake/model", dtype="fp32",
        )
