aixtools 0.1.10__py3-none-any.whl → 0.2.0__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 (62) hide show
  1. aixtools/_version.py +2 -2
  2. aixtools/agents/agent.py +26 -7
  3. aixtools/agents/print_nodes.py +54 -0
  4. aixtools/agents/prompt.py +2 -2
  5. aixtools/compliance/private_data.py +1 -1
  6. aixtools/evals/discovery.py +174 -0
  7. aixtools/evals/evals.py +74 -0
  8. aixtools/evals/run_evals.py +110 -0
  9. aixtools/logging/log_objects.py +24 -23
  10. aixtools/mcp/client.py +148 -2
  11. aixtools/server/__init__.py +0 -6
  12. aixtools/server/path.py +88 -31
  13. aixtools/testing/aix_test_model.py +9 -1
  14. aixtools/tools/doctor/mcp_tool_doctor.py +79 -0
  15. aixtools/tools/doctor/tool_doctor.py +4 -0
  16. aixtools/tools/doctor/tool_recommendation.py +5 -0
  17. aixtools/utils/config.py +0 -1
  18. {aixtools-0.1.10.dist-info → aixtools-0.2.0.dist-info}/METADATA +186 -30
  19. {aixtools-0.1.10.dist-info → aixtools-0.2.0.dist-info}/RECORD +23 -55
  20. aixtools-0.2.0.dist-info/entry_points.txt +4 -0
  21. aixtools-0.2.0.dist-info/top_level.txt +1 -0
  22. aixtools/server/workspace_privacy.py +0 -65
  23. aixtools-0.1.10.dist-info/entry_points.txt +0 -2
  24. aixtools-0.1.10.dist-info/top_level.txt +0 -5
  25. docker/mcp-base/Dockerfile +0 -33
  26. docker/mcp-base/zscaler.crt +0 -28
  27. notebooks/example_faulty_mcp_server.ipynb +0 -74
  28. notebooks/example_mcp_server_stdio.ipynb +0 -76
  29. notebooks/example_raw_mcp_client.ipynb +0 -84
  30. notebooks/example_tool_doctor.ipynb +0 -65
  31. scripts/config.sh +0 -28
  32. scripts/lint.sh +0 -32
  33. scripts/log_view.sh +0 -18
  34. scripts/run_example_mcp_server.sh +0 -14
  35. scripts/run_faulty_mcp_server.sh +0 -13
  36. scripts/run_server.sh +0 -29
  37. scripts/test.sh +0 -30
  38. tests/unit/__init__.py +0 -0
  39. tests/unit/a2a/__init__.py +0 -0
  40. tests/unit/a2a/google_sdk/__init__.py +0 -0
  41. tests/unit/a2a/google_sdk/pydantic_ai_adapter/__init__.py +0 -0
  42. tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_agent_executor.py +0 -188
  43. tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_storage.py +0 -156
  44. tests/unit/a2a/google_sdk/test_card.py +0 -114
  45. tests/unit/a2a/google_sdk/test_remote_agent_connection.py +0 -413
  46. tests/unit/a2a/google_sdk/test_utils.py +0 -208
  47. tests/unit/agents/__init__.py +0 -0
  48. tests/unit/agents/test_prompt.py +0 -363
  49. tests/unit/compliance/test_private_data.py +0 -329
  50. tests/unit/google/__init__.py +0 -1
  51. tests/unit/google/test_client.py +0 -233
  52. tests/unit/mcp/__init__.py +0 -0
  53. tests/unit/mcp/test_client.py +0 -242
  54. tests/unit/server/__init__.py +0 -0
  55. tests/unit/server/test_path.py +0 -225
  56. tests/unit/server/test_utils.py +0 -362
  57. tests/unit/utils/__init__.py +0 -0
  58. tests/unit/utils/test_files.py +0 -146
  59. tests/unit/vault/__init__.py +0 -0
  60. tests/unit/vault/test_vault.py +0 -246
  61. {tests → aixtools/evals}/__init__.py +0 -0
  62. {aixtools-0.1.10.dist-info → aixtools-0.2.0.dist-info}/WHEEL +0 -0
@@ -1,329 +0,0 @@
1
- import json
2
- import unittest
3
- from pathlib import Path
4
-
5
- from aixtools.compliance.private_data import PrivateData, PRIVATE_DATA_FILE
6
- from aixtools.server.path import get_workspace_path
7
-
8
-
9
- class TestPrivateData(unittest.TestCase):
10
- """Test cases for PrivateData class without mocking or patching."""
11
-
12
- def setUp(self):
13
- """Set up test environment with temporary directory."""
14
- # Create a test user and session context
15
- self.workspace_path = Path(get_workspace_path())
16
- self.workspace_path.mkdir(parents=True, exist_ok=True)
17
- self.private_data_file = self.workspace_path / PRIVATE_DATA_FILE
18
-
19
- def tearDown(self):
20
- """Clean up test environment."""
21
- self.private_data_file.unlink(missing_ok=True)
22
-
23
- def test_init_no_existing_file(self):
24
- """Test initialization when no private data file exists."""
25
- private_data = PrivateData()
26
- self.assertFalse(private_data.has_private_data)
27
- self.assertEqual(private_data.get_private_datasets(), [])
28
- self.assertEqual(private_data.get_idap_datasets(), [])
29
-
30
- def test_init_with_existing_file(self):
31
- """Test initialization when private data file exists."""
32
- # Create a test file
33
- private_data_path = self.workspace_path / PRIVATE_DATA_FILE
34
- test_data = {
35
- "_has_private_data": True,
36
- "_private_datasets": ["dataset1", "dataset2"],
37
- "_idap_datasets": ["dataset1"],
38
- "ctx": None
39
- }
40
- with open(private_data_path, "w") as f:
41
- json.dump(test_data, f)
42
-
43
- private_data = PrivateData()
44
-
45
- self.assertTrue(private_data.has_private_data)
46
- self.assertEqual(set(private_data.get_private_datasets()), {"dataset1", "dataset2"})
47
- self.assertEqual(private_data.get_idap_datasets(), ["dataset1"])
48
-
49
- def test_add_private_dataset_new(self):
50
- """Test adding a new private dataset."""
51
- private_data = PrivateData()
52
-
53
- private_data.add_private_dataset("test_dataset")
54
-
55
- self.assertTrue(private_data.has_private_data)
56
- self.assertIn("test_dataset", private_data.get_private_datasets())
57
-
58
- # Verify data is saved to file
59
- private_data_path = self.workspace_path / PRIVATE_DATA_FILE
60
- self.assertTrue(private_data_path.exists())
61
-
62
- with open(private_data_path, "r") as f:
63
- saved_data = json.load(f)
64
-
65
- self.assertTrue(saved_data["_has_private_data"])
66
- self.assertIn("test_dataset", saved_data["_private_datasets"])
67
-
68
- def test_add_private_dataset_duplicate(self):
69
- """Test adding a duplicate private dataset."""
70
- private_data = PrivateData()
71
-
72
- private_data.add_private_dataset("test_dataset")
73
- initial_count = len(private_data.get_private_datasets())
74
-
75
- private_data.add_private_dataset("test_dataset")
76
- final_count = len(private_data.get_private_datasets())
77
-
78
- self.assertEqual(initial_count, final_count)
79
- self.assertEqual(private_data.get_private_datasets().count("test_dataset"), 1)
80
-
81
- def test_add_idap_dataset_new(self):
82
- """Test adding a new IDAP dataset."""
83
- private_data = PrivateData()
84
-
85
- private_data.add_idap_dataset("idap_dataset")
86
-
87
- self.assertTrue(private_data.has_private_data)
88
- self.assertIn("idap_dataset", private_data.get_idap_datasets())
89
- self.assertIn("idap_dataset", private_data.get_private_datasets())
90
-
91
- # Verify data is saved to file
92
- private_data_path = self.workspace_path / PRIVATE_DATA_FILE
93
- self.assertTrue(private_data_path.exists())
94
-
95
- with open(private_data_path, "r") as f:
96
- saved_data = json.load(f)
97
-
98
- self.assertTrue(saved_data["_has_private_data"])
99
- self.assertIn("idap_dataset", saved_data["_idap_datasets"])
100
- self.assertIn("idap_dataset", saved_data["_private_datasets"])
101
-
102
- def test_add_idap_dataset_existing_private(self):
103
- """Test adding IDAP dataset when it already exists as private dataset."""
104
- private_data = PrivateData()
105
-
106
- private_data.add_private_dataset("existing_dataset")
107
- private_data.add_idap_dataset("existing_dataset")
108
-
109
- self.assertTrue(private_data.has_private_data)
110
- self.assertIn("existing_dataset", private_data.get_idap_datasets())
111
- self.assertIn("existing_dataset", private_data.get_private_datasets())
112
- self.assertEqual(private_data.get_private_datasets().count("existing_dataset"), 1)
113
-
114
- def test_add_idap_dataset_duplicate(self):
115
- """Test adding a duplicate IDAP dataset."""
116
- private_data = PrivateData()
117
-
118
- private_data.add_idap_dataset("idap_dataset")
119
- initial_idap_count = len(private_data.get_idap_datasets())
120
- initial_private_count = len(private_data.get_private_datasets())
121
-
122
- private_data.add_idap_dataset("idap_dataset")
123
- final_idap_count = len(private_data.get_idap_datasets())
124
- final_private_count = len(private_data.get_private_datasets())
125
-
126
- self.assertEqual(initial_idap_count, final_idap_count)
127
- self.assertEqual(initial_private_count, final_private_count)
128
-
129
- def test_has_private_dataset(self):
130
- """Test checking if a private dataset exists."""
131
- private_data = PrivateData()
132
-
133
- self.assertFalse(private_data.has_private_dataset("nonexistent"))
134
-
135
- private_data.add_private_dataset("test_dataset")
136
-
137
- self.assertTrue(private_data.has_private_dataset("test_dataset"))
138
- self.assertFalse(private_data.has_private_dataset("nonexistent"))
139
-
140
- def test_has_idap_dataset(self):
141
- """Test checking if an IDAP dataset exists."""
142
- private_data = PrivateData()
143
-
144
- self.assertFalse(private_data.has_idap_dataset("nonexistent"))
145
-
146
- private_data.add_idap_dataset("idap_dataset")
147
-
148
- self.assertTrue(private_data.has_idap_dataset("idap_dataset"))
149
- self.assertFalse(private_data.has_idap_dataset("nonexistent"))
150
-
151
- def test_get_datasets_returns_copies(self):
152
- """Test that get methods return copies to prevent external modification."""
153
- private_data = PrivateData()
154
- private_data.add_private_dataset("private_dataset")
155
- private_data.add_idap_dataset("idap_dataset")
156
-
157
- private_datasets = private_data.get_private_datasets()
158
- idap_datasets = private_data.get_idap_datasets()
159
-
160
- # Modify the returned lists
161
- private_datasets.append("external_modification")
162
- idap_datasets.append("external_modification")
163
-
164
- # Verify original data is unchanged
165
- self.assertNotIn("external_modification", private_data.get_private_datasets())
166
- self.assertNotIn("external_modification", private_data.get_idap_datasets())
167
-
168
- def test_has_private_data_setter_true(self):
169
- """Test setting has_private_data to True."""
170
- private_data = PrivateData()
171
-
172
- private_data.has_private_data = True
173
-
174
- self.assertTrue(private_data.has_private_data)
175
-
176
- # Verify data is saved to file
177
- private_data_path = self.workspace_path / PRIVATE_DATA_FILE
178
- self.assertTrue(private_data_path.exists())
179
-
180
- def test_has_private_data_setter_false(self):
181
- """Test setting has_private_data to False clears all data."""
182
- private_data = PrivateData()
183
- private_data.add_private_dataset("dataset1")
184
- private_data.add_idap_dataset("dataset2")
185
-
186
- private_data.has_private_data = False
187
-
188
- self.assertFalse(private_data.has_private_data)
189
- self.assertEqual(private_data.get_private_datasets(), [])
190
- self.assertEqual(private_data.get_idap_datasets(), [])
191
-
192
- # Verify file is deleted
193
- private_data_path = self.workspace_path / PRIVATE_DATA_FILE
194
- self.assertFalse(private_data_path.exists())
195
-
196
- def test_save_with_no_private_data(self):
197
- """Test saving when has_private_data is False deletes the file."""
198
- private_data = PrivateData()
199
- private_data.add_private_dataset("test_dataset")
200
-
201
- # Verify file exists
202
- private_data_path = self.workspace_path / PRIVATE_DATA_FILE
203
- self.assertTrue(private_data_path.exists())
204
-
205
- private_data.has_private_data = False
206
-
207
- # Verify file is deleted
208
- self.assertFalse(private_data_path.exists())
209
-
210
- def test_save_with_private_data(self):
211
- """Test saving when has_private_data is True creates/updates the file."""
212
- private_data = PrivateData()
213
- private_data.add_private_dataset("dataset1")
214
- private_data.add_idap_dataset("dataset2")
215
-
216
- private_data_path = self.workspace_path / PRIVATE_DATA_FILE
217
- self.assertTrue(private_data_path.exists())
218
-
219
- with open(private_data_path, "r") as f:
220
- saved_data = json.load(f)
221
-
222
- self.assertTrue(saved_data["_has_private_data"])
223
- self.assertIn("dataset1", saved_data["_private_datasets"])
224
- self.assertIn("dataset2", saved_data["_private_datasets"])
225
- self.assertIn("dataset2", saved_data["_idap_datasets"])
226
- self.assertIsNone(saved_data["ctx"])
227
-
228
- def test_load_from_corrupted_file(self):
229
- """Test loading from a corrupted JSON file."""
230
- private_data_path = self.workspace_path / PRIVATE_DATA_FILE
231
- with open(private_data_path, "w") as f:
232
- f.write("invalid json content")
233
-
234
- with self.assertRaises(json.JSONDecodeError):
235
- PrivateData()
236
-
237
- def test_load_from_file_with_missing_fields(self):
238
- """Test loading from a file with missing fields uses defaults."""
239
- private_data_path = self.workspace_path / PRIVATE_DATA_FILE
240
- test_data = {"_has_private_data": True} # Missing other fields
241
- with open(private_data_path, "w") as f:
242
- json.dump(test_data, f)
243
-
244
- private_data = PrivateData()
245
-
246
- self.assertTrue(private_data.has_private_data)
247
- self.assertEqual(private_data.get_private_datasets(), [])
248
- self.assertEqual(private_data.get_idap_datasets(), [])
249
-
250
- def test_multiple_operations_persistence(self):
251
- """Test that multiple operations are properly persisted."""
252
- # Create first instance and add data
253
- private_data1 = PrivateData()
254
- private_data1.add_private_dataset("dataset1")
255
- private_data1.add_idap_dataset("dataset2")
256
-
257
- # Create second instance and verify data is loaded
258
- private_data2 = PrivateData()
259
- self.assertTrue(private_data2.has_private_data)
260
- self.assertIn("dataset1", private_data2.get_private_datasets())
261
- self.assertIn("dataset2", private_data2.get_private_datasets())
262
- self.assertIn("dataset2", private_data2.get_idap_datasets())
263
-
264
- # Add more data with second instance
265
- private_data2.add_private_dataset("dataset3")
266
-
267
- # Create third instance and verify all data is present
268
- private_data3 = PrivateData()
269
- expected_private = {"dataset1", "dataset2", "dataset3"}
270
- expected_idap = {"dataset2"}
271
-
272
- self.assertEqual(set(private_data3.get_private_datasets()), expected_private)
273
- self.assertEqual(set(private_data3.get_idap_datasets()), expected_idap)
274
-
275
- def test_str_and_repr(self):
276
- """Test string representation methods."""
277
- private_data = PrivateData()
278
- private_data.add_private_dataset("dataset1")
279
- private_data.add_idap_dataset("dataset2")
280
-
281
- str_repr = str(private_data)
282
- repr_repr = repr(private_data)
283
-
284
- self.assertEqual(str_repr, repr_repr)
285
- self.assertIn("has_private_data=True", str_repr)
286
- self.assertIn("dataset1", str_repr)
287
- self.assertIn("dataset2", str_repr)
288
- self.assertIn(str(self.workspace_path / PRIVATE_DATA_FILE), str_repr)
289
-
290
- def test_file_operations_create_directories(self):
291
- """Test that file operations create necessary directories."""
292
- # Remove the workspace directory
293
- import shutil
294
- shutil.rmtree(self.workspace_path)
295
- self.assertFalse(self.workspace_path.exists())
296
-
297
- # Create PrivateData instance and add data
298
- private_data = PrivateData()
299
- private_data.add_private_dataset("test_dataset")
300
-
301
- # Verify directory and file were created
302
- self.assertTrue(self.workspace_path.exists())
303
- private_data_path = self.workspace_path / PRIVATE_DATA_FILE
304
- self.assertTrue(private_data_path.exists())
305
-
306
- def test_concurrent_modifications_simulation(self):
307
- """Test simulation of concurrent modifications (without actual threading)."""
308
- # Create two instances with same context
309
- private_data1 = PrivateData()
310
- private_data2 = PrivateData()
311
-
312
- # Add data with first instance
313
- private_data1.add_private_dataset("dataset1")
314
-
315
- # Second instance should see the changes when reloaded
316
- private_data2.load()
317
- self.assertIn("dataset1", private_data2.get_private_datasets())
318
-
319
- # Add data with second instance
320
- private_data2.add_idap_dataset("dataset2")
321
-
322
- # First instance should see the changes when reloaded
323
- private_data1.load()
324
- self.assertIn("dataset2", private_data1.get_idap_datasets())
325
- self.assertIn("dataset2", private_data1.get_private_datasets())
326
-
327
-
328
- if __name__ == "__main__":
329
- unittest.main()
@@ -1 +0,0 @@
1
- """Unit tests for aixtools.google module."""
@@ -1,233 +0,0 @@
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