LLM/DocsGPT
1
0
Fork 0
mirror of https://github.com/arc53/DocsGPT.git synced 2025-11-29 16:43:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

add redis mock to anthropic and sagemaker

Browse Source
This commit is contained in:
fadingNA
2024-10-14 11:54:22 -04:00
parent 3e32724729
commit adb2cf35d4
2 changed files with 62 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
tests/llm/test_sagemaker.py
View File

@@ -52,28 +52,38 @@ class TestSagemakerAPILLM(unittest.TestCase):
self.response['Body'].read.return_value.decode.return_value = json.dumps(self.result)
def test_gen(self):
with patch.object(self.sagemaker.runtime, 'invoke_endpoint',
return_value=self.response) as mock_invoke_endpoint:
output = self.sagemaker.gen(None, self.messages)
mock_invoke_endpoint.assert_called_once_with(
EndpointName=self.sagemaker.endpoint,
ContentType='application/json',
Body=self.body_bytes
)
self.assertEqual(output,
self.result[0]['generated_text'][len(self.prompt):])
with patch('application.cache.make_redis') as mock_make_redis:
mock_redis_instance = mock_make_redis.return_value
mock_redis_instance.get.return_value = None
with patch.object(self.sagemaker.runtime, 'invoke_endpoint',
return_value=self.response) as mock_invoke_endpoint:
output = self.sagemaker.gen(None, self.messages)
mock_invoke_endpoint.assert_called_once_with(
EndpointName=self.sagemaker.endpoint,
ContentType='application/json',
Body=self.body_bytes
)
self.assertEqual(output,
self.result[0]['generated_text'][len(self.prompt):])
mock_make_redis.assert_called_once()
mock_redis_instance.set.assert_called_once()
def test_gen_stream(self):
with patch.object(self.sagemaker.runtime, 'invoke_endpoint_with_response_stream',
return_value=self.response) as mock_invoke_endpoint:
output = list(self.sagemaker.gen_stream(None, self.messages))
mock_invoke_endpoint.assert_called_once_with(
EndpointName=self.sagemaker.endpoint,
ContentType='application/json',
Body=self.body_bytes_stream
)
self.assertEqual(output, [])
with patch('application.cache.make_redis') as mock_make_redis:
mock_redis_instance = mock_make_redis.return_value
mock_redis_instance.get.return_value = None
with patch.object(self.sagemaker.runtime, 'invoke_endpoint_with_response_stream',
return_value=self.response) as mock_invoke_endpoint:
output = list(self.sagemaker.gen_stream(None, self.messages))
mock_invoke_endpoint.assert_called_once_with(
EndpointName=self.sagemaker.endpoint,
ContentType='application/json',
Body=self.body_bytes_stream
)
self.assertEqual(output, [])
mock_redis_instance.set.assert_called_once()
class TestLineIterator(unittest.TestCase):
def setUp(self):