Former-commit-id: 4fe62f34ed
dockerize
Kye 1 year ago
parent 9f658b6d31
commit 74ac5af137

@ -1,13 +1,13 @@
import pytest
import json import json
from unittest.mock import MagicMock
from unittest.mock import patch
import os import os
from swarms.structs.flow import Flow, stop_when_repeats from unittest.mock import MagicMock, patch
from swarms.models import OpenAIChat
import pytest
from dotenv import load_dotenv from dotenv import load_dotenv
from swarms.models import OpenAIChat
from swarms.structs.flow import Flow, stop_when_repeats
load_dotenv() load_dotenv()
openai_api_key = os.getenv("OPENAI_API_KEY") openai_api_key = os.getenv("OPENAI_API_KEY")
@ -89,7 +89,8 @@ def test_env_variable_handling(monkeypatch):
# Test initializing the flow with different stopping conditions # Test initializing the flow with different stopping conditions
def test_flow_with_custom_stopping_condition(mocked_llm): def test_flow_with_custom_stopping_condition(mocked_llm):
stopping_condition = lambda x: "terminate" in x.lower() def stopping_condition(x):
return "terminate" in x.lower()
flow = Flow(llm=mocked_llm, max_loops=5, stopping_condition=stopping_condition) flow = Flow(llm=mocked_llm, max_loops=5, stopping_condition=stopping_condition)
assert flow.stopping_condition("Please terminate now") assert flow.stopping_condition("Please terminate now")
assert not flow.stopping_condition("Continue the process") assert not flow.stopping_condition("Continue the process")
@ -165,7 +166,7 @@ def test_mocked_openai_chat(MockedOpenAIChat):
llm = MockedOpenAIChat(openai_api_key=openai_api_key) llm = MockedOpenAIChat(openai_api_key=openai_api_key)
llm.return_value = MagicMock() llm.return_value = MagicMock()
flow = Flow(llm=llm, max_loops=5) flow = Flow(llm=llm, max_loops=5)
response = flow.run("Mocked run") flow.run("Mocked run")
assert MockedOpenAIChat.called assert MockedOpenAIChat.called
# Test retry attempts # Test retry attempts

Loading…
Cancel
Save