feat: tests for generator

Author: Darinochka

tests/generation/utterances/test_basic_synthesizer.py

@@ -1,7 +1,5 @@
 from unittest.mock import AsyncMock, Mock
 
-import pytest
-
 from autointent.generation.utterances import SynthesizerChatTemplate, UtteranceGenerator
 
 
@@ -53,23 +51,6 @@ def test_on_dataset(dataset):
     assert n_before + len(new_samples) == n_after
     assert len(new_samples) == len(dataset.intents)
 
-@pytest.mark.asyncio
-async def test_default_chat_template_async(dataset):
-    template = SynthesizerChatTemplate(dataset, split="train_0")
-    prompt = template(dataset.intents[0], n_examples=1)
-    for msg in prompt:
-        assert not has_unfilled_fields(msg["content"])
-    assert "extra_instructions" not in prompt
-
-
-@pytest.mark.asyncio
-async def test_extra_instructions_async(dataset):
-    template = SynthesizerChatTemplate(dataset, split="train_0", extra_instructions="football")
-    prompt = template(dataset.intents[0], n_examples=1)[0]["content"]
-    assert "extra_instructions" not in prompt
-    assert "football" in prompt
-
-
 def test_on_dataset_async(dataset):
     mock_llm = AsyncMock()
     mock_llm.get_chat_completion_async.return_value = "1. LLM answer"

tests/generation/utterances/test_evolver.py

@@ -1,4 +1,4 @@
-from unittest.mock import Mock
+from unittest.mock import AsyncMock, Mock
 
 from autointent.generation.utterances import AbstractEvolution, UtteranceEvolver
 
@@ -32,3 +32,27 @@ def test_on_dataset(dataset):
 
     assert n_before + len(new_samples) == n_after
     assert len(new_samples) == n_before
+
+def test_on_dataset_evolver_async(dataset):
+    mock_llm = AsyncMock()
+    mock_llm.get_chat_completion_async.return_value = "LLM answer"
+
+    split_name = "train_0"
+
+    template = AbstractEvolution()
+    augmenter = UtteranceEvolver(mock_llm, [template], async_mode=True)
+
+    n_before = len(dataset[split_name])
+    new_samples = augmenter.augment(dataset, split_name=split_name, n_evolutions=1, update_split=False)
+    n_after = len(dataset[split_name])
+
+    assert n_before == n_after
+    assert len(new_samples) == n_before
+    assert all(sample.utterance == "LLM answer" for sample in new_samples)
+
+    n_before = len(dataset[split_name])
+    new_samples = augmenter.augment(dataset, split_name=split_name, n_evolutions=1, update_split=True)
+    n_after = len(dataset[split_name])
+
+    assert n_before + len(new_samples) == n_after
+    assert len(new_samples) == n_before

tests/generation/utterances/test_generator.py

@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
@@ -36,3 +36,30 @@ def test_get_chat_completion(mock_openai_client):
 
     assert response == "Test response"
     mock_openai_client.return_value.chat.completions.create.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_get_chat_completion_async():
+    test_messages = [Message(role="user", content="Hello, how are you?")]
+
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock(message=MagicMock(content="I'm fine, thank you!"))]
+
+    with patch("openai.AsyncOpenAI") as mock_async_openai:
+        mock_instance = mock_async_openai.return_value
+        mock_instance.chat.completions.create = AsyncMock(return_value=mock_response)
+
+        generator = Generator()
+
+        result = await generator.get_chat_completion_async(test_messages)
+
+        assert result == "I'm fine, thank you!"
+
+        mock_instance.chat.completions.create.assert_awaited_once_with(
+            messages=test_messages,
+            model=generator.model_name,
+            max_tokens=150,
+            n=1,
+            stop=None,
+            temperature=0.7,
+        )