dify/api/tests/integration_tests/model_runtime/__mock/google.py

from collections.abc import Generator
from unittest.mock import MagicMock

import google.generativeai.types.generation_types as generation_config_types
import pytest
from _pytest.monkeypatch import MonkeyPatch
from google.ai import generativelanguage as glm
from google.ai.generativelanguage_v1beta.types import content as gag_content
from google.generativeai import GenerativeModel
from google.generativeai.types import GenerateContentResponse, content_types, safety_types
from google.generativeai.types.generation_types import BaseGenerateContentResponse

from extensions import ext_redis


class MockGoogleResponseClass:
    _done = False

    def __iter__(self):
        full_response_text = "it's google!"

        for i in range(0, len(full_response_text) + 1, 1):
            if i == len(full_response_text):
                self._done = True
                yield GenerateContentResponse(
                    done=True, iterator=None, result=glm.GenerateContentResponse({}), chunks=[]
                )
            else:
                yield GenerateContentResponse(
                    done=False, iterator=None, result=glm.GenerateContentResponse({}), chunks=[]
                )


class MockGoogleResponseCandidateClass:
    finish_reason = "stop"

    @property
    def content(self) -> gag_content.Content:
        return gag_content.Content(parts=[gag_content.Part(text="it's google!")])


class MockGoogleClass:
    @staticmethod
    def generate_content_sync() -> GenerateContentResponse:
        return GenerateContentResponse(done=True, iterator=None, result=glm.GenerateContentResponse({}), chunks=[])

    @staticmethod
    def generate_content_stream() -> Generator[GenerateContentResponse, None, None]:
        return MockGoogleResponseClass()

    def generate_content(
        self: GenerativeModel,
        contents: content_types.ContentsType,
        *,
        generation_config: generation_config_types.GenerationConfigType | None = None,
        safety_settings: safety_types.SafetySettingOptions | None = None,
        stream: bool = False,
        **kwargs,
    ) -> GenerateContentResponse:
        if stream:
            return MockGoogleClass.generate_content_stream()

        return MockGoogleClass.generate_content_sync()

    @property
    def generative_response_text(self) -> str:
        return "it's google!"

    @property
    def generative_response_candidates(self) -> list[MockGoogleResponseCandidateClass]:
        return [MockGoogleResponseCandidateClass()]


def mock_configure(api_key: str):
    if len(api_key) < 16:
        raise Exception("Invalid API key")


class MockFileState:
    def __init__(self):
        self.name = "FINISHED"


class MockGoogleFile:
    def __init__(self, name: str = "mock_file_name"):
        self.name = name
        self.state = MockFileState()


def mock_get_file(name: str) -> MockGoogleFile:
    return MockGoogleFile(name)


def mock_upload_file(path: str, mime_type: str) -> MockGoogleFile:
    return MockGoogleFile()


@pytest.fixture
def setup_google_mock(request, monkeypatch: MonkeyPatch):
    monkeypatch.setattr(BaseGenerateContentResponse, "text", MockGoogleClass.generative_response_text)
    monkeypatch.setattr(BaseGenerateContentResponse, "candidates", MockGoogleClass.generative_response_candidates)
    monkeypatch.setattr(GenerativeModel, "generate_content", MockGoogleClass.generate_content)
    monkeypatch.setattr("google.generativeai.configure", mock_configure)
    monkeypatch.setattr("google.generativeai.get_file", mock_get_file)
    monkeypatch.setattr("google.generativeai.upload_file", mock_upload_file)

    yield

    monkeypatch.undo()


@pytest.fixture
def setup_mock_redis() -> None:
    ext_redis.redis_client.get = MagicMock(return_value=None)
    ext_redis.redis_client.setex = MagicMock(return_value=None)
    ext_redis.redis_client.exists = MagicMock(return_value=True)