aixtools 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of aixtools might be problematic. Click here for more details.

Files changed (46) hide show
  1. aixtools/_version.py +2 -2
  2. aixtools/a2a/app.py +1 -1
  3. aixtools/a2a/google_sdk/__init__.py +0 -0
  4. aixtools/a2a/google_sdk/card.py +27 -0
  5. aixtools/a2a/google_sdk/pydantic_ai_adapter/agent_executor.py +199 -0
  6. aixtools/a2a/google_sdk/pydantic_ai_adapter/storage.py +26 -0
  7. aixtools/a2a/google_sdk/remote_agent_connection.py +88 -0
  8. aixtools/a2a/google_sdk/utils.py +59 -0
  9. aixtools/agents/prompt.py +97 -0
  10. aixtools/context.py +5 -0
  11. aixtools/google/client.py +25 -0
  12. aixtools/logging/logging_config.py +45 -0
  13. aixtools/mcp/client.py +274 -0
  14. aixtools/mcp/faulty_mcp.py +7 -7
  15. aixtools/server/utils.py +3 -3
  16. aixtools/utils/config.py +6 -0
  17. aixtools/utils/files.py +17 -0
  18. aixtools/utils/utils.py +7 -0
  19. {aixtools-0.1.3.dist-info → aixtools-0.1.5.dist-info}/METADATA +3 -1
  20. {aixtools-0.1.3.dist-info → aixtools-0.1.5.dist-info}/RECORD +45 -14
  21. {aixtools-0.1.3.dist-info → aixtools-0.1.5.dist-info}/top_level.txt +1 -0
  22. scripts/test.sh +23 -0
  23. tests/__init__.py +0 -0
  24. tests/unit/__init__.py +0 -0
  25. tests/unit/a2a/__init__.py +0 -0
  26. tests/unit/a2a/google_sdk/__init__.py +0 -0
  27. tests/unit/a2a/google_sdk/pydantic_ai_adapter/__init__.py +0 -0
  28. tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_agent_executor.py +188 -0
  29. tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_storage.py +156 -0
  30. tests/unit/a2a/google_sdk/test_card.py +114 -0
  31. tests/unit/a2a/google_sdk/test_remote_agent_connection.py +413 -0
  32. tests/unit/a2a/google_sdk/test_utils.py +208 -0
  33. tests/unit/agents/__init__.py +0 -0
  34. tests/unit/agents/test_prompt.py +363 -0
  35. tests/unit/google/__init__.py +1 -0
  36. tests/unit/google/test_client.py +233 -0
  37. tests/unit/mcp/__init__.py +0 -0
  38. tests/unit/mcp/test_client.py +242 -0
  39. tests/unit/server/__init__.py +0 -0
  40. tests/unit/server/test_path.py +225 -0
  41. tests/unit/server/test_utils.py +362 -0
  42. tests/unit/utils/__init__.py +0 -0
  43. tests/unit/utils/test_files.py +146 -0
  44. aixtools/a2a/__init__.py +0 -5
  45. {aixtools-0.1.3.dist-info → aixtools-0.1.5.dist-info}/WHEEL +0 -0
  46. {aixtools-0.1.3.dist-info → aixtools-0.1.5.dist-info}/entry_points.txt +0 -0
@@ -0,0 +1,363 @@
1
+ """Unit tests for aixtools.agents.prompt module."""
2
+
3
+ import tempfile
4
+ import unittest
5
+ from pathlib import Path
6
+ from unittest.mock import Mock, patch, mock_open
7
+
8
+ from pydantic_ai import BinaryContent
9
+
10
+ from aixtools.agents.prompt import (
11
+ CLAUDE_IMAGE_MAX_FILE_SIZE_IN_CONTEXT,
12
+ CLAUDE_MAX_FILE_SIZE_IN_CONTEXT,
13
+ build_user_input,
14
+ file_to_binary_content,
15
+ should_be_included_into_context,
16
+ )
17
+
18
+
19
+ class TestShouldBeIncludedIntoContext(unittest.TestCase):
20
+ """Test cases for should_be_included_into_context function."""
21
+
22
+ def test_non_binary_content_returns_false(self):
23
+ """Test that non-BinaryContent returns False."""
24
+ result = should_be_included_into_context("text content", 100)
25
+ self.assertFalse(result)
26
+
27
+ result = should_be_included_into_context(None, 100)
28
+ self.assertFalse(result)
29
+
30
+ def test_text_media_type_returns_false(self):
31
+ """Test that text media types return False."""
32
+ binary_content = BinaryContent(data=b"test", media_type="text/plain")
33
+ result = should_be_included_into_context(binary_content, 100)
34
+ self.assertFalse(result)
35
+
36
+ binary_content = BinaryContent(data=b"test", media_type="text/html")
37
+ result = should_be_included_into_context(binary_content, 100)
38
+ self.assertFalse(result)
39
+
40
+ def test_archive_types_return_false(self):
41
+ """Test that archive media types return False."""
42
+ archive_types = [
43
+ "application/zip",
44
+ "application/x-tar",
45
+ "application/gzip",
46
+ "application/x-gzip",
47
+ "application/x-rar-compressed",
48
+ "application/x-7z-compressed",
49
+ ]
50
+
51
+ for media_type in archive_types:
52
+ binary_content = BinaryContent(data=b"test", media_type=media_type)
53
+ result = should_be_included_into_context(binary_content, 100)
54
+ self.assertFalse(result, f"Archive type {media_type} should return False")
55
+
56
+ def test_image_within_size_limit_returns_true(self):
57
+ """Test that images within size limit return True."""
58
+ with patch.object(BinaryContent, 'is_image', new_callable=lambda: True):
59
+ binary_content = BinaryContent(data=b"image_data", media_type="image/png")
60
+
61
+ # Test with size under limit
62
+ result = should_be_included_into_context(binary_content, 1024)
63
+ self.assertTrue(result)
64
+
65
+ def test_image_over_size_limit_returns_false(self):
66
+ """Test that images over size limit return False."""
67
+ with patch.object(BinaryContent, 'is_image', new_callable=lambda: True):
68
+ binary_content = BinaryContent(data=b"image_data", media_type="image/png")
69
+
70
+ # Test with size over limit
71
+ result = should_be_included_into_context(binary_content, CLAUDE_IMAGE_MAX_FILE_SIZE_IN_CONTEXT + 1)
72
+ self.assertFalse(result)
73
+
74
+ def test_non_image_within_size_limit_returns_true(self):
75
+ """Test that non-images within size limit return True."""
76
+ with patch.object(BinaryContent, 'is_image', new_callable=lambda: False):
77
+ binary_content = BinaryContent(data=b"pdf_data", media_type="application/pdf")
78
+
79
+ # Test with size under limit
80
+ result = should_be_included_into_context(binary_content, 1024)
81
+ self.assertTrue(result)
82
+
83
+ def test_non_image_over_size_limit_returns_false(self):
84
+ """Test that non-images over size limit return False."""
85
+ with patch.object(BinaryContent, 'is_image', new_callable=lambda: False):
86
+ binary_content = BinaryContent(data=b"pdf_data", media_type="application/pdf")
87
+
88
+ # Test with size over limit
89
+ result = should_be_included_into_context(binary_content, CLAUDE_MAX_FILE_SIZE_IN_CONTEXT + 1)
90
+ self.assertFalse(result)
91
+
92
+ def test_custom_size_limits(self):
93
+ """Test with custom size limits."""
94
+ # Test non-image content with custom limits (since image detection is complex)
95
+ pdf_content = BinaryContent(data=b"pdf_data", media_type="application/pdf")
96
+
97
+ # Test non-image over custom file limit
98
+ result = should_be_included_into_context(
99
+ pdf_content, 5000, max_img_size_bytes=1024, max_file_size_bytes=4096
100
+ )
101
+ self.assertFalse(result)
102
+
103
+ # Test non-image under custom file limit
104
+ result = should_be_included_into_context(
105
+ pdf_content, 2000, max_img_size_bytes=1024, max_file_size_bytes=4096
106
+ )
107
+ self.assertTrue(result)
108
+
109
+ # Test with very small custom limits
110
+ result = should_be_included_into_context(
111
+ pdf_content, 100, max_img_size_bytes=50, max_file_size_bytes=80
112
+ )
113
+ self.assertFalse(result)
114
+
115
+ # Test with very large custom limits
116
+ result = should_be_included_into_context(
117
+ pdf_content, 100, max_img_size_bytes=200, max_file_size_bytes=200
118
+ )
119
+ self.assertTrue(result)
120
+
121
+
122
+ class TestFileToBinaryContent(unittest.TestCase):
123
+ """Test cases for file_to_binary_content function."""
124
+
125
+ def setUp(self):
126
+ """Set up test fixtures."""
127
+ self.temp_dir = tempfile.mkdtemp()
128
+ self.temp_path = Path(self.temp_dir)
129
+
130
+ def tearDown(self):
131
+ """Clean up test fixtures."""
132
+ import shutil
133
+ shutil.rmtree(self.temp_dir)
134
+
135
+ @patch('aixtools.agents.prompt.is_text_content')
136
+ @patch('mimetypes.guess_type')
137
+ def test_text_file_returns_string(self, mock_guess_type, mock_is_text):
138
+ """Test that text files return decoded strings."""
139
+ mock_guess_type.return_value = ("text/plain", None)
140
+ mock_is_text.return_value = True
141
+
142
+ test_file = self.temp_path / "test.txt"
143
+ test_content = "Hello, world!"
144
+ test_file.write_text(test_content, encoding="utf-8")
145
+
146
+ result = file_to_binary_content(test_file)
147
+
148
+ self.assertEqual(result, test_content)
149
+ mock_is_text.assert_called_once()
150
+
151
+ @patch('aixtools.agents.prompt.is_text_content')
152
+ @patch('mimetypes.guess_type')
153
+ def test_binary_file_returns_binary_content(self, mock_guess_type, mock_is_text):
154
+ """Test that binary files return BinaryContent."""
155
+ mock_guess_type.return_value = ("image/png", None)
156
+ mock_is_text.return_value = False
157
+
158
+ test_file = self.temp_path / "test.png"
159
+ test_data = b'\x89PNG\r\n\x1a\n'
160
+ test_file.write_bytes(test_data)
161
+
162
+ result = file_to_binary_content(test_file)
163
+
164
+ self.assertIsInstance(result, BinaryContent)
165
+ if isinstance(result, BinaryContent):
166
+ self.assertEqual(result.data, test_data)
167
+ self.assertEqual(result.media_type, "image/png")
168
+
169
+ @patch('aixtools.agents.prompt.is_text_content')
170
+ @patch('mimetypes.guess_type')
171
+ def test_unknown_mime_type_defaults_to_octet_stream(self, mock_guess_type, mock_is_text):
172
+ """Test that unknown mime types default to application/octet-stream."""
173
+ mock_guess_type.return_value = (None, None)
174
+ mock_is_text.return_value = False
175
+
176
+ test_file = self.temp_path / "test.unknown"
177
+ test_data = b'unknown data'
178
+ test_file.write_bytes(test_data)
179
+
180
+ result = file_to_binary_content(test_file)
181
+
182
+ self.assertIsInstance(result, BinaryContent)
183
+ if isinstance(result, BinaryContent):
184
+ self.assertEqual(result.media_type, "application/octet-stream")
185
+
186
+ @patch('aixtools.agents.prompt.is_text_content')
187
+ def test_explicit_mime_type(self, mock_is_text):
188
+ """Test with explicitly provided mime type."""
189
+ mock_is_text.return_value = False
190
+
191
+ test_file = self.temp_path / "test.data"
192
+ test_data = b'test data'
193
+ test_file.write_bytes(test_data)
194
+
195
+ result = file_to_binary_content(test_file, "application/custom")
196
+
197
+ self.assertIsInstance(result, BinaryContent)
198
+ if isinstance(result, BinaryContent):
199
+ self.assertEqual(result.media_type, "application/custom")
200
+
201
+
202
+ class TestBuildUserInput(unittest.TestCase):
203
+ """Test cases for build_user_input function."""
204
+
205
+ def setUp(self):
206
+ """Set up test fixtures."""
207
+ self.session_tuple = ("user123", "session456")
208
+ self.user_text = "Please analyze these files"
209
+
210
+ def test_no_file_paths_returns_text_only(self):
211
+ """Test that no file paths returns just the user text."""
212
+ result = build_user_input(self.session_tuple, self.user_text, [])
213
+
214
+ self.assertEqual(result, self.user_text)
215
+
216
+ @patch('aixtools.agents.prompt.should_be_included_into_context')
217
+ @patch('aixtools.agents.prompt.file_to_binary_content')
218
+ @patch('aixtools.agents.prompt.container_to_host_path')
219
+ @patch('mimetypes.guess_type')
220
+ def test_single_file_not_included_in_context(
221
+ self, mock_guess_type, mock_container_to_host,
222
+ mock_file_to_binary, mock_should_include
223
+ ):
224
+ """Test with single file that should not be included in context."""
225
+ # Setup mocks
226
+ mock_guess_type.return_value = ("text/plain", None)
227
+ mock_should_include.return_value = False
228
+
229
+ mock_host_path = Mock()
230
+ mock_host_path.stat.return_value.st_size = 1024
231
+ mock_container_to_host.return_value = mock_host_path
232
+
233
+ mock_file_to_binary.return_value = "file content"
234
+
235
+ file_paths = [Path("/workspace/test.txt")]
236
+
237
+ result = build_user_input(self.session_tuple, self.user_text, file_paths)
238
+
239
+ self.assertIsInstance(result, list)
240
+ self.assertEqual(len(result), 1) # Only the prompt, no binary attachments
241
+ self.assertIsInstance(result[0], str)
242
+ prompt_text = str(result[0])
243
+ self.assertIn("Please analyze these files", prompt_text)
244
+ self.assertIn("Attachments:", prompt_text)
245
+ self.assertIn("test.txt", prompt_text)
246
+ self.assertIn("file_size=1024 bytes", prompt_text)
247
+
248
+ @patch('aixtools.agents.prompt.should_be_included_into_context')
249
+ @patch('aixtools.agents.prompt.file_to_binary_content')
250
+ @patch('aixtools.agents.prompt.container_to_host_path')
251
+ @patch('mimetypes.guess_type')
252
+ def test_single_file_included_in_context(
253
+ self, mock_guess_type, mock_container_to_host,
254
+ mock_file_to_binary, mock_should_include
255
+ ):
256
+ """Test with single file that should be included in context."""
257
+ # Setup mocks
258
+ mock_guess_type.return_value = ("image/png", None)
259
+ mock_should_include.return_value = True
260
+
261
+ mock_host_path = Mock()
262
+ mock_host_path.stat.return_value.st_size = 2048
263
+ mock_container_to_host.return_value = mock_host_path
264
+
265
+ mock_binary_content = BinaryContent(data=b"image data", media_type="image/png")
266
+ mock_file_to_binary.return_value = mock_binary_content
267
+
268
+ file_paths = [Path("/workspace/image.png")]
269
+
270
+ result = build_user_input(self.session_tuple, self.user_text, file_paths)
271
+
272
+ self.assertIsInstance(result, list)
273
+ self.assertEqual(len(result), 2) # Prompt + 1 binary attachment
274
+ self.assertIsInstance(result[0], str)
275
+ prompt_text = str(result[0])
276
+ self.assertIn("Please analyze these files", prompt_text)
277
+ self.assertIn("Attachments:", prompt_text)
278
+ self.assertIn("image.png", prompt_text)
279
+ self.assertIn("file_size=2048 bytes", prompt_text)
280
+ self.assertIn("provided to model context at index 0", prompt_text)
281
+ self.assertEqual(result[1], mock_binary_content)
282
+
283
+ @patch('aixtools.agents.prompt.should_be_included_into_context')
284
+ @patch('aixtools.agents.prompt.file_to_binary_content')
285
+ @patch('aixtools.agents.prompt.container_to_host_path')
286
+ @patch('mimetypes.guess_type')
287
+ def test_multiple_files_mixed_inclusion(
288
+ self, mock_guess_type, mock_container_to_host,
289
+ mock_file_to_binary, mock_should_include
290
+ ):
291
+ """Test with multiple files, some included in context, some not."""
292
+ # Setup mocks
293
+ mock_guess_type.side_effect = [("text/plain", None), ("image/png", None)]
294
+ mock_should_include.side_effect = [False, True] # First file not included, second included
295
+
296
+ mock_host_path1 = Mock()
297
+ mock_host_path1.stat.return_value.st_size = 1024
298
+ mock_host_path2 = Mock()
299
+ mock_host_path2.stat.return_value.st_size = 2048
300
+ mock_container_to_host.side_effect = [mock_host_path1, mock_host_path2]
301
+
302
+ mock_text_content = "text content"
303
+ mock_binary_content = BinaryContent(data=b"image data", media_type="image/png")
304
+ mock_file_to_binary.side_effect = [mock_text_content, mock_binary_content]
305
+
306
+ file_paths = [Path("/workspace/text.txt"), Path("/workspace/image.png")]
307
+
308
+ result = build_user_input(self.session_tuple, self.user_text, file_paths)
309
+
310
+ self.assertIsInstance(result, list)
311
+ self.assertEqual(len(result), 2) # Prompt + 1 binary attachment
312
+ self.assertIsInstance(result[0], str)
313
+ prompt_text = str(result[0])
314
+ self.assertIn("Please analyze these files", prompt_text)
315
+ self.assertIn("Attachments:", prompt_text)
316
+ self.assertIn("text.txt", prompt_text)
317
+ self.assertIn("image.png", prompt_text)
318
+ self.assertIn("file_size=1024 bytes", prompt_text)
319
+ self.assertIn("file_size=2048 bytes", prompt_text)
320
+ self.assertIn("provided to model context at index 0", prompt_text)
321
+ self.assertEqual(result[1], mock_binary_content)
322
+
323
+ def test_non_workspace_path_raises_error(self):
324
+ """Test that non-workspace paths raise ValueError."""
325
+ file_paths = [Path("/invalid/path.txt")]
326
+
327
+ with self.assertRaises(ValueError) as context:
328
+ build_user_input(self.session_tuple, self.user_text, file_paths)
329
+
330
+ self.assertIn(
331
+ "Container path must be a subdir of '/workspace', got '/invalid/path.txt' instead",
332
+ str(context.exception),
333
+ )
334
+
335
+ @patch('aixtools.agents.prompt.should_be_included_into_context')
336
+ @patch('aixtools.agents.prompt.file_to_binary_content')
337
+ @patch('aixtools.agents.prompt.container_to_host_path')
338
+ @patch('mimetypes.guess_type')
339
+ def test_unknown_mime_type_defaults(
340
+ self, mock_guess_type, mock_container_to_host,
341
+ mock_file_to_binary, mock_should_include
342
+ ):
343
+ """Test that unknown mime types default to application/octet-stream."""
344
+ # Setup mocks
345
+ mock_guess_type.return_value = (None, None) # Unknown mime type
346
+ mock_should_include.return_value = False
347
+
348
+ mock_host_path = Mock()
349
+ mock_host_path.stat.return_value.st_size = 1024
350
+ mock_container_to_host.return_value = mock_host_path
351
+
352
+ mock_file_to_binary.return_value = "file content"
353
+
354
+ file_paths = [Path("/workspace/unknown.dat")]
355
+
356
+ build_user_input(self.session_tuple, self.user_text, file_paths)
357
+
358
+ # Verify that file_to_binary_content was called with the default mime type
359
+ mock_file_to_binary.assert_called_once_with(mock_host_path, "application/octet-stream")
360
+
361
+
362
+ if __name__ == '__main__':
363
+ unittest.main()
@@ -0,0 +1 @@
1
+ """Unit tests for aixtools.google module."""
@@ -0,0 +1,233 @@
1
+ """Unit tests for aixtools.google.client module."""
2
+
3
+ import os
4
+ import unittest
5
+ from pathlib import Path
6
+ from unittest.mock import Mock, patch, MagicMock
7
+
8
+ from aixtools.google.client import get_genai_client
9
+
10
+
11
+ class TestGetGenaiClient(unittest.TestCase):
12
+ """Test cases for get_genai_client function."""
13
+
14
+ def setUp(self):
15
+ """Set up test fixtures."""
16
+ # Store original environment variables to restore later
17
+ self.original_env = os.environ.copy()
18
+
19
+ def tearDown(self):
20
+ """Clean up after tests."""
21
+ # Restore original environment variables
22
+ os.environ.clear()
23
+ os.environ.update(self.original_env)
24
+
25
+ @patch('aixtools.google.client.genai.Client')
26
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', 'test-project')
27
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', 'us-central1')
28
+ @patch('aixtools.google.client.GOOGLE_GENAI_USE_VERTEXAI', True)
29
+ def test_get_genai_client_basic_success(self, mock_client):
30
+ """Test get_genai_client with basic valid configuration."""
31
+ mock_client_instance = Mock()
32
+ mock_client.return_value = mock_client_instance
33
+
34
+ result = get_genai_client()
35
+
36
+ mock_client.assert_called_once_with(
37
+ vertexai=True,
38
+ project='test-project',
39
+ location='us-central1'
40
+ )
41
+ self.assertEqual(result, mock_client_instance)
42
+
43
+ @patch('aixtools.google.client.genai.Client')
44
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', 'test-project')
45
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', 'us-west1')
46
+ @patch('aixtools.google.client.GOOGLE_GENAI_USE_VERTEXAI', False)
47
+ def test_get_genai_client_without_vertexai(self, mock_client):
48
+ """Test get_genai_client with GOOGLE_GENAI_USE_VERTEXAI set to False."""
49
+ mock_client_instance = Mock()
50
+ mock_client.return_value = mock_client_instance
51
+
52
+ result = get_genai_client()
53
+
54
+ mock_client.assert_called_once_with(
55
+ vertexai=False,
56
+ project='test-project',
57
+ location='us-west1'
58
+ )
59
+ self.assertEqual(result, mock_client_instance)
60
+
61
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', None)
62
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', 'us-central1')
63
+ def test_get_genai_client_missing_project(self):
64
+ """Test get_genai_client raises AssertionError when GOOGLE_CLOUD_PROJECT is not set."""
65
+ with self.assertRaises(AssertionError) as context:
66
+ get_genai_client()
67
+
68
+ self.assertIn("GOOGLE_CLOUD_PROJECT is not set", str(context.exception))
69
+
70
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', 'test-project')
71
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', None)
72
+ def test_get_genai_client_missing_location(self):
73
+ """Test get_genai_client raises AssertionError when GOOGLE_CLOUD_LOCATION is not set."""
74
+ with self.assertRaises(AssertionError) as context:
75
+ get_genai_client()
76
+
77
+ self.assertIn("GOOGLE_CLOUD_LOCATION is not set", str(context.exception))
78
+
79
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', '')
80
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', 'us-central1')
81
+ def test_get_genai_client_empty_project(self):
82
+ """Test get_genai_client raises AssertionError when GOOGLE_CLOUD_PROJECT is empty."""
83
+ with self.assertRaises(AssertionError) as context:
84
+ get_genai_client()
85
+
86
+ self.assertIn("GOOGLE_CLOUD_PROJECT is not set", str(context.exception))
87
+
88
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', 'test-project')
89
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', '')
90
+ def test_get_genai_client_empty_location(self):
91
+ """Test get_genai_client raises AssertionError when GOOGLE_CLOUD_LOCATION is empty."""
92
+ with self.assertRaises(AssertionError) as context:
93
+ get_genai_client()
94
+
95
+ self.assertIn("GOOGLE_CLOUD_LOCATION is not set", str(context.exception))
96
+
97
+ @patch('aixtools.google.client.genai.Client')
98
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', 'test-project')
99
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', 'us-central1')
100
+ @patch('aixtools.google.client.GOOGLE_GENAI_USE_VERTEXAI', True)
101
+ @patch('aixtools.google.client.logger')
102
+ def test_get_genai_client_with_valid_service_account_key(self, mock_logger, mock_client):
103
+ """Test get_genai_client with valid service account key file."""
104
+ mock_client_instance = Mock()
105
+ mock_client.return_value = mock_client_instance
106
+
107
+ # Create a temporary file path that exists
108
+ with patch('pathlib.Path.exists', return_value=True):
109
+ service_account_path = Path('/tmp/test-service-account.json')
110
+
111
+ result = get_genai_client(service_account_path)
112
+
113
+ # Verify the environment variable was set
114
+ self.assertEqual(os.environ.get('GOOGLE_APPLICATION_CREDENTIALS'), str(service_account_path))
115
+
116
+ # Verify logging was called
117
+ mock_logger.info.assert_called_once_with(f"✅ GCP Service Account Key File: {service_account_path}")
118
+
119
+ # Verify client was created correctly
120
+ mock_client.assert_called_once_with(
121
+ vertexai=True,
122
+ project='test-project',
123
+ location='us-central1'
124
+ )
125
+ self.assertEqual(result, mock_client_instance)
126
+
127
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', 'test-project')
128
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', 'us-central1')
129
+ def test_get_genai_client_with_nonexistent_service_account_key(self):
130
+ """Test get_genai_client raises FileNotFoundError for nonexistent service account key."""
131
+ with patch('pathlib.Path.exists', return_value=False):
132
+ service_account_path = Path('/tmp/nonexistent-service-account.json')
133
+
134
+ with self.assertRaises(FileNotFoundError) as context:
135
+ get_genai_client(service_account_path)
136
+
137
+ self.assertIn(f"Service account key file not found: {service_account_path}", str(context.exception))
138
+
139
+ @patch('aixtools.google.client.genai.Client')
140
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', 'test-project')
141
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', 'us-central1')
142
+ @patch('aixtools.google.client.GOOGLE_GENAI_USE_VERTEXAI', True)
143
+ def test_get_genai_client_preserves_existing_credentials_env(self, mock_client):
144
+ """Test get_genai_client preserves existing GOOGLE_APPLICATION_CREDENTIALS when no service account key provided."""
145
+ mock_client_instance = Mock()
146
+ mock_client.return_value = mock_client_instance
147
+
148
+ # Set existing credentials
149
+ existing_credentials = '/existing/path/to/credentials.json'
150
+ os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = existing_credentials
151
+
152
+ result = get_genai_client()
153
+
154
+ # Verify existing credentials are preserved
155
+ self.assertEqual(os.environ.get('GOOGLE_APPLICATION_CREDENTIALS'), existing_credentials)
156
+
157
+ mock_client.assert_called_once_with(
158
+ vertexai=True,
159
+ project='test-project',
160
+ location='us-central1'
161
+ )
162
+ self.assertEqual(result, mock_client_instance)
163
+
164
+ @patch('aixtools.google.client.genai.Client')
165
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', 'test-project')
166
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', 'us-central1')
167
+ @patch('aixtools.google.client.GOOGLE_GENAI_USE_VERTEXAI', True)
168
+ @patch('aixtools.google.client.logger')
169
+ def test_get_genai_client_overwrites_existing_credentials_env(self, mock_logger, mock_client):
170
+ """Test get_genai_client overwrites existing GOOGLE_APPLICATION_CREDENTIALS when service account key provided."""
171
+ mock_client_instance = Mock()
172
+ mock_client.return_value = mock_client_instance
173
+
174
+ # Set existing credentials
175
+ existing_credentials = '/existing/path/to/credentials.json'
176
+ os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = existing_credentials
177
+
178
+ with patch('pathlib.Path.exists', return_value=True):
179
+ new_service_account_path = Path('/tmp/new-service-account.json')
180
+
181
+ result = get_genai_client(new_service_account_path)
182
+
183
+ # Verify new credentials overwrite existing ones
184
+ self.assertEqual(os.environ.get('GOOGLE_APPLICATION_CREDENTIALS'), str(new_service_account_path))
185
+
186
+ # Verify logging was called
187
+ mock_logger.info.assert_called_once_with(f"✅ GCP Service Account Key File: {new_service_account_path}")
188
+
189
+ mock_client.assert_called_once_with(
190
+ vertexai=True,
191
+ project='test-project',
192
+ location='us-central1'
193
+ )
194
+ self.assertEqual(result, mock_client_instance)
195
+
196
+ @patch('aixtools.google.client.genai.Client')
197
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', 'test-project')
198
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', 'us-central1')
199
+ @patch('aixtools.google.client.GOOGLE_GENAI_USE_VERTEXAI', True)
200
+ def test_get_genai_client_genai_client_exception(self, mock_client):
201
+ """Test get_genai_client propagates exceptions from genai.Client."""
202
+ mock_client.side_effect = Exception("GenAI client initialization failed")
203
+
204
+ with self.assertRaises(Exception) as context:
205
+ get_genai_client()
206
+
207
+ self.assertIn("GenAI client initialization failed", str(context.exception))
208
+
209
+ @patch('aixtools.google.client.genai.Client')
210
+ @patch('aixtools.google.client.GOOGLE_CLOUD_PROJECT', 'test-project')
211
+ @patch('aixtools.google.client.GOOGLE_CLOUD_LOCATION', 'us-central1')
212
+ @patch('aixtools.google.client.GOOGLE_GENAI_USE_VERTEXAI', True)
213
+ @patch('aixtools.google.client.logger')
214
+ def test_get_genai_client_service_account_path_as_string(self, mock_logger, mock_client):
215
+ """Test get_genai_client handles service account path as string correctly."""
216
+ mock_client_instance = Mock()
217
+ mock_client.return_value = mock_client_instance
218
+
219
+ with patch('pathlib.Path.exists', return_value=True):
220
+ service_account_path = Path('/tmp/test-service-account.json')
221
+
222
+ result = get_genai_client(service_account_path)
223
+
224
+ # Verify the environment variable was set as string
225
+ self.assertEqual(os.environ.get('GOOGLE_APPLICATION_CREDENTIALS'), str(service_account_path))
226
+ self.assertIsInstance(os.environ.get('GOOGLE_APPLICATION_CREDENTIALS'), str)
227
+
228
+ mock_client.assert_called_once()
229
+ self.assertEqual(result, mock_client_instance)
230
+
231
+
232
+ if __name__ == '__main__':
233
+ unittest.main()
File without changes