Underground-Digital/Workflow-Engine

1import os2from collections.abc import Callable3from typing import Any, Literal4 5import pytest6 7# import monkeypatch8from _pytest.monkeypatch import MonkeyPatch9from nomic import embed10 11 12def create_embedding(texts: list[str], model: str, **kwargs: Any) -> dict:13    texts_len = len(texts)14 15    foo_embedding_sample = 0.12345616 17    combined = {18        "embeddings": [[foo_embedding_sample for _ in range(768)] for _ in range(texts_len)],19        "usage": {"prompt_tokens": texts_len, "total_tokens": texts_len},20        "model": model,21        "inference_mode": "remote",22    }23 24    return combined25 26 27def mock_nomic(28    monkeypatch: MonkeyPatch,29    methods: list[Literal["text_embedding"]],30) -> Callable[[], None]:31    """32    mock nomic module33 34    :param monkeypatch: pytest monkeypatch fixture35    :return: unpatch function36    """37 38    def unpatch() -> None:39        monkeypatch.undo()40 41    if "text_embedding" in methods:42        monkeypatch.setattr(embed, "text", create_embedding)43 44    return unpatch45 46 47MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"48 49 50@pytest.fixture51def setup_nomic_mock(request, monkeypatch):52    methods = request.param if hasattr(request, "param") else []53    if MOCK:54        unpatch = mock_nomic(monkeypatch, methods=methods)55 56    yield57 58    if MOCK:59        unpatch()60