refactor:
- changed method get_agent_tools for only get_tools on agent interface - changed OllamaModelManager to correc semantic OllamaService test: - added test for service initialization - added tests to list tools and create tools into PDF agent
This commit is contained in:
0
test/service/ollama/__init__.py
Normal file
0
test/service/ollama/__init__.py
Normal file
84
test/service/ollama/test_ollamaService.py
Normal file
84
test/service/ollama/test_ollamaService.py
Normal file
@@ -0,0 +1,84 @@
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
from src.service.ollama.ollamaService import OllamaService
|
||||
|
||||
|
||||
|
||||
|
||||
class testOllamaService(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.model_manager = OllamaService()
|
||||
self.model_manager.base_url = "http://test-url:11434"
|
||||
|
||||
@patch('src.service.ollama.ollamaService.requests.get')
|
||||
def test_get_available_models_success(self, mock_get):
|
||||
# Setup mock response
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
'models': [
|
||||
{'name': 'model1'},
|
||||
{'name': 'model2'}
|
||||
]
|
||||
}
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
# Call method
|
||||
result = self.model_manager.get_available_models()
|
||||
|
||||
# Assertions
|
||||
mock_get.assert_called_once_with("http://test-url:11434/api/tags")
|
||||
self.assertEqual(len(result), 2)
|
||||
self.assertEqual(result, [{'name': 'model1'}, {'name': 'model2'}])
|
||||
|
||||
@patch('src.service.ollama.ollamaService.requests.get')
|
||||
def test_get_available_models_exception(self, mock_get):
|
||||
# Setup mock to raise exception
|
||||
mock_get.side_effect = Exception("Connection error")
|
||||
|
||||
# Call method
|
||||
result = self.model_manager.get_available_models()
|
||||
|
||||
# Assertions
|
||||
self.assertEqual(result, [])
|
||||
|
||||
@patch('src.service.ollama.ollamaService.requests.post')
|
||||
def test_get_model_details_success(self, mock_post):
|
||||
# Setup mock response
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = {
|
||||
'parameters': {'context_length': 4096},
|
||||
'modelfile': {'parameter': 'llama2'},
|
||||
'license': 'Apache 2.0'
|
||||
}
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
# Call method
|
||||
result = self.model_manager.get_model_details('llama2')
|
||||
|
||||
# Assertions
|
||||
mock_post.assert_called_once_with(
|
||||
"http://test-url:11434/api/show",
|
||||
json={"name": "llama2"}
|
||||
)
|
||||
self.assertEqual(result['name'], 'llama2')
|
||||
self.assertEqual(result['context_size'], 4096)
|
||||
self.assertEqual(result['license'], 'Apache 2.0')
|
||||
|
||||
@patch('src.service.ollama.ollamaService.requests.post')
|
||||
def test_get_model_details_exception(self, mock_post):
|
||||
# Setup mock to raise exception
|
||||
mock_post.side_effect = Exception("API error")
|
||||
|
||||
# Call method
|
||||
result = self.model_manager.get_model_details('unknown_model')
|
||||
|
||||
# Assertions
|
||||
self.assertEqual(result['name'], 'unknown_model')
|
||||
self.assertTrue('error' in result)
|
||||
self.assertEqual(result['error'], 'API error')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user