Merge pull request #595 from better629/feat_gemini

Feat gemini

tests/metagpt/provider/test_google_gemini_api.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+# @Desc   : the unittest of google gemini api
+
+from abc import ABC
+from dataclasses import dataclass
+
+import pytest
+
+from metagpt.provider.google_gemini_api import GeminiGPTAPI
+
+messages = [{"role": "user", "content": "who are you"}]
+
+
+@dataclass
+class MockGeminiResponse(ABC):
+    text: str
+
+
+default_resp = MockGeminiResponse(text="I'm gemini from google")
+
+
+def mock_llm_ask(self, messages: list[dict]) -> MockGeminiResponse:
+    return default_resp
+
+
+def test_gemini_completion(mocker):
+    mocker.patch("metagpt.provider.google_gemini_api.GeminiGPTAPI.completion", mock_llm_ask)
+    resp = GeminiGPTAPI().completion(messages)
+    assert resp.text == default_resp.text
+
+
+async def mock_llm_aask(self, messgaes: list[dict]) -> MockGeminiResponse:
+    return default_resp
+
+
+@pytest.mark.asyncio
+async def test_gemini_acompletion(mocker):
+    mocker.patch("metagpt.provider.google_gemini_api.GeminiGPTAPI.acompletion", mock_llm_aask)
+    resp = await GeminiGPTAPI().acompletion(messages)
+    assert resp.text == default_resp.text