Skip to content

Commit a229a5e

Browse files
Deepthi-ChandDeepthi-Chand
committed
fix tests
1 parent 6c69ce3 commit a229a5e

File tree

3 files changed

+108
-110
lines changed

3 files changed

+108
-110
lines changed

‎tests/test_aimodels.py‎

Lines changed: 40 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@ def test_init(self) -> None:
2020
self.assertEqual(self.client.base_url, self.base_url)
2121
self.assertEqual(self.client.auth_client, self.auth_client)
2222

23-
@patch.object(AIModelClient, "get")
24-
def test_search_models(self, mock_get: MagicMock) -> None:
23+
@patch.object(AIModelClient, "_make_request")
24+
def test_search_models(self, mock_request: MagicMock) -> None:
2525
"""Test AI model search."""
26-
mock_get.return_value = {
26+
mock_request.return_value = {
2727
"total": 5,
2828
"results": [{"id": "1", "displayName": "Test Model", "modelType": "LLM"}],
2929
}
@@ -33,12 +33,12 @@ def test_search_models(self, mock_get: MagicMock) -> None:
3333
self.assertEqual(result["total"], 5)
3434
self.assertEqual(len(result["results"]), 1)
3535
self.assertEqual(result["results"][0]["displayName"], "Test Model")
36-
mock_get.assert_called_once()
36+
mock_request.assert_called_once()
3737

38-
@patch.object(AIModelClient, "get")
39-
def test_get_model_by_id(self, mock_get: MagicMock) -> None:
38+
@patch.object(AIModelClient, "_make_request")
39+
def test_get_model_by_id(self, mock_request: MagicMock) -> None:
4040
"""Test get AI model by ID."""
41-
mock_get.return_value = {
41+
mock_request.return_value = {
4242
"id": "123",
4343
"displayName": "Test Model",
4444
"modelType": "LLM",
@@ -90,37 +90,37 @@ def test_get_organization_models(self, mock_post: MagicMock) -> None:
9090
self.assertIsInstance(result, (list, dict))
9191
mock_post.assert_called_once()
9292

93-
@patch.object(AIModelClient, "get")
94-
def test_search_with_filters(self, mock_get: MagicMock) -> None:
93+
@patch.object(AIModelClient, "_make_request")
94+
def test_search_with_filters(self, mock_request: MagicMock) -> None:
9595
"""Test AI model search with filters."""
96-
mock_get.return_value = {"total": 3, "results": []}
96+
mock_request.return_value = {"total": 3, "results": []}
9797

9898
result = self.client.search(
9999
query="language",
100100
tags=["nlp"],
101101
sectors=["tech"],
102+
status="ACTIVE",
102103
model_type="LLM",
103104
provider="OpenAI",
104-
status="ACTIVE",
105105
)
106106

107107
self.assertEqual(result["total"], 3)
108-
mock_get.assert_called_once()
108+
mock_request.assert_called_once()
109109

110-
@patch.object(AIModelClient, "post")
111-
def test_graphql_error_handling(self, mock_post: MagicMock) -> None:
110+
@patch.object(AIModelClient, "_make_request")
111+
def test_graphql_error_handling(self, mock_request: MagicMock) -> None:
112112
"""Test GraphQL error handling."""
113113
from dataspace_sdk.exceptions import DataSpaceAPIError
114114

115-
mock_post.return_value = {"errors": [{"message": "GraphQL error"}]}
115+
mock_request.return_value = {"errors": [{"message": "GraphQL error"}]}
116116

117117
with self.assertRaises(DataSpaceAPIError):
118118
self.client.get_by_id_graphql("123")
119119

120-
@patch.object(AIModelClient, "post")
121-
def test_call_model(self, mock_post: MagicMock) -> None:
120+
@patch.object(AIModelClient, "_make_request")
121+
def test_call_model(self, mock_request: MagicMock) -> None:
122122
"""Test calling an AI model."""
123-
mock_post.return_value = {
123+
mock_request.return_value = {
124124
"success": True,
125125
"output": "Paris is the capital of France.",
126126
"latency_ms": 150,
@@ -136,12 +136,12 @@ def test_call_model(self, mock_post: MagicMock) -> None:
136136
self.assertTrue(result["success"])
137137
self.assertEqual(result["output"], "Paris is the capital of France.")
138138
self.assertEqual(result["latency_ms"], 150)
139-
mock_post.assert_called_once()
139+
mock_request.assert_called_once()
140140

141-
@patch.object(AIModelClient, "post")
142-
def test_call_model_async(self, mock_post: MagicMock) -> None:
141+
@patch.object(AIModelClient, "_make_request")
142+
def test_call_model_async(self, mock_request: MagicMock) -> None:
143143
"""Test calling an AI model asynchronously."""
144-
mock_post.return_value = {
144+
mock_request.return_value = {
145145
"task_id": "task-456",
146146
"status": "PENDING",
147147
"created_at": "2024-01-01T00:00:00Z",
@@ -155,12 +155,12 @@ def test_call_model_async(self, mock_post: MagicMock) -> None:
155155

156156
self.assertEqual(result["task_id"], "task-456")
157157
self.assertEqual(result["status"], "PENDING")
158-
mock_post.assert_called_once()
158+
mock_request.assert_called_once()
159159

160-
@patch.object(AIModelClient, "post")
161-
def test_call_model_error(self, mock_post: MagicMock) -> None:
160+
@patch.object(AIModelClient, "_make_request")
161+
def test_call_model_error(self, mock_request: MagicMock) -> None:
162162
"""Test AI model call with error."""
163-
mock_post.return_value = {
163+
mock_request.return_value = {
164164
"success": False,
165165
"error": "Model not available",
166166
}
@@ -173,10 +173,10 @@ def test_call_model_error(self, mock_post: MagicMock) -> None:
173173
self.assertFalse(result["success"])
174174
self.assertEqual(result["error"], "Model not available")
175175

176-
@patch.object(AIModelClient, "post")
177-
def test_create_model(self, mock_post: MagicMock) -> None:
176+
@patch.object(AIModelClient, "_make_request")
177+
def test_create_model(self, mock_request: MagicMock) -> None:
178178
"""Test creating an AI model."""
179-
mock_post.return_value = {
179+
mock_request.return_value = {
180180
"id": "new-model-123",
181181
"displayName": "New Model",
182182
"modelType": "LLM",
@@ -192,36 +192,33 @@ def test_create_model(self, mock_post: MagicMock) -> None:
192192

193193
self.assertEqual(result["id"], "new-model-123")
194194
self.assertEqual(result["displayName"], "New Model")
195-
mock_post.assert_called_once()
195+
mock_request.assert_called_once()
196196

197-
@patch.object(AIModelClient, "put")
198-
def test_update_model(self, mock_put: MagicMock) -> None:
197+
@patch.object(AIModelClient, "_make_request")
198+
def test_update_model(self, mock_request: MagicMock) -> None:
199199
"""Test updating an AI model."""
200-
mock_put.return_value = {
200+
mock_request.return_value = {
201201
"id": "123",
202202
"displayName": "Updated Model",
203-
"modelType": "LLM",
203+
"description": "Updated description",
204204
}
205205

206206
result = self.client.update(
207-
"123",
208-
{
209-
"displayName": "Updated Model",
210-
},
207+
"123", {"displayName": "Updated Model", "description": "Updated description"}
211208
)
212209

213210
self.assertEqual(result["displayName"], "Updated Model")
214-
mock_put.assert_called_once()
211+
mock_request.assert_called_once()
215212

216-
@patch.object(AIModelClient, "delete")
217-
def test_delete_model(self, mock_delete: MagicMock) -> None:
213+
@patch.object(AIModelClient, "_make_request")
214+
def test_delete_model(self, mock_request: MagicMock) -> None:
218215
"""Test deleting an AI model."""
219-
mock_delete.return_value = {"message": "Model deleted successfully"}
216+
mock_request.return_value = {"message": "Model deleted successfully"}
220217

221218
result = self.client.delete_model("123")
222219

223220
self.assertEqual(result["message"], "Model deleted successfully")
224-
mock_delete.assert_called_once()
221+
mock_request.assert_called_once()
225222

226223

227224
if __name__ == "__main__":

‎tests/test_datasets.py‎

Lines changed: 51 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@ def test_init(self) -> None:
2020
self.assertEqual(self.client.base_url, self.base_url)
2121
self.assertEqual(self.client.auth_client, self.auth_client)
2222

23-
@patch.object(DatasetClient, "get")
24-
def test_search_datasets(self, mock_get: MagicMock) -> None:
23+
@patch.object(DatasetClient, "_make_request")
24+
def test_search_datasets(self, mock_request: MagicMock) -> None:
2525
"""Test dataset search."""
26-
mock_get.return_value = {
26+
mock_request.return_value = {
2727
"total": 10,
2828
"results": [{"id": "1", "title": "Test Dataset"}],
2929
}
@@ -32,22 +32,22 @@ def test_search_datasets(self, mock_get: MagicMock) -> None:
3232

3333
self.assertEqual(result["total"], 10)
3434
self.assertEqual(len(result["results"]), 1)
35-
mock_get.assert_called_once()
35+
mock_request.assert_called_once()
3636

37-
@patch.object(DatasetClient, "post")
38-
def test_get_dataset_by_id(self, mock_post: MagicMock) -> None:
37+
@patch.object(DatasetClient, "_make_request")
38+
def test_get_dataset_by_id(self, mock_request: MagicMock) -> None:
3939
"""Test get dataset by ID."""
40-
mock_post.return_value = {"data": {"dataset": {"id": "123", "title": "Test Dataset"}}}
40+
mock_request.return_value = {"data": {"dataset": {"id": "123", "title": "Test Dataset"}}}
4141

4242
result = self.client.get_by_id("123")
4343

4444
self.assertEqual(result["id"], "123")
4545
self.assertEqual(result["title"], "Test Dataset")
4646

47-
@patch.object(DatasetClient, "post")
48-
def test_list_all_datasets(self, mock_post: MagicMock) -> None:
47+
@patch.object(DatasetClient, "_make_request")
48+
def test_list_all_datasets(self, mock_request: MagicMock) -> None:
4949
"""Test list all datasets."""
50-
mock_post.return_value = {"data": {"datasets": [{"id": "1", "title": "Dataset 1"}]}}
50+
mock_request.return_value = {"data": {"datasets": [{"id": "1", "title": "Dataset 1"}]}}
5151

5252
result = self.client.list_all(limit=10, offset=0)
5353

@@ -73,10 +73,10 @@ def test_get_organization_datasets(self, mock_post: MagicMock) -> None:
7373
self.assertIsInstance(result, (list, dict))
7474
mock_post.assert_called_once()
7575

76-
@patch.object(DatasetClient, "get")
77-
def test_search_with_filters(self, mock_get: MagicMock) -> None:
76+
@patch.object(DatasetClient, "_make_request")
77+
def test_search_with_filters(self, mock_request: MagicMock) -> None:
7878
"""Test dataset search with filters."""
79-
mock_get.return_value = {"total": 5, "results": []}
79+
mock_request.return_value = {"total": 5, "results": []}
8080

8181
result = self.client.search(
8282
query="health",
@@ -87,52 +87,57 @@ def test_search_with_filters(self, mock_get: MagicMock) -> None:
8787
)
8888

8989
self.assertEqual(result["total"], 5)
90-
mock_get.assert_called_once()
90+
mock_request.assert_called_once()
9191

92-
@patch.object(DatasetClient, "get")
93-
def test_get_resources(self, mock_get: MagicMock) -> None:
94-
"""Test get dataset resources."""
95-
mock_get.return_value = [
96-
{
97-
"id": "res-1",
98-
"title": "Resource 1",
99-
"format": "CSV",
100-
"url": "https://example.com/data.csv",
92+
@patch.object(DatasetClient, "_make_request")
93+
def test_get_dataset_with_resources(self, mock_request: MagicMock) -> None:
94+
"""Test get dataset by ID which includes resources."""
95+
mock_request.return_value = {
96+
"data": {
97+
"dataset": {
98+
"id": "dataset-123",
99+
"title": "Test Dataset",
100+
"resources": [
101+
{
102+
"id": "res-1",
103+
"title": "Resource 1",
104+
"fileDetails": {"format": "CSV"},
105+
}
106+
],
107+
}
101108
}
102-
]
109+
}
103110

104-
result = self.client.get_resources("dataset-123")
111+
result = self.client.get_by_id("dataset-123")
105112

106-
self.assertEqual(len(result), 1)
107-
self.assertEqual(result[0]["title"], "Resource 1")
108-
mock_get.assert_called_once()
113+
self.assertEqual(result["id"], "dataset-123")
114+
self.assertEqual(len(result["resources"]), 1)
115+
self.assertEqual(result["resources"][0]["title"], "Resource 1")
116+
mock_request.assert_called_once()
109117

110-
@patch.object(DatasetClient, "post")
111-
def test_list_by_organization(self, mock_post: MagicMock) -> None:
112-
"""Test list datasets by organization."""
113-
mock_post.return_value = {
114-
"data": {
115-
"datasets": [
116-
{"id": "1", "title": "Org Dataset 1"},
117-
{"id": "2", "title": "Org Dataset 2"},
118-
]
119-
}
120-
}
118+
@patch.object(DatasetClient, "list_all")
119+
def test_get_organization_datasets(self, mock_list_all: MagicMock) -> None:
120+
"""Test get datasets by organization."""
121+
mock_list_all.return_value = [
122+
{"id": "1", "title": "Org Dataset 1"},
123+
{"id": "2", "title": "Org Dataset 2"},
124+
]
121125

122-
result = self.client.list_by_organization("org-123", limit=10)
126+
result = self.client.get_organization_datasets("org-123", limit=10)
123127

124-
self.assertIsInstance(result, (list, dict))
125-
mock_post.assert_called_once()
128+
self.assertIsInstance(result, list)
129+
self.assertEqual(len(result), 2)
130+
mock_list_all.assert_called_once_with(organization_id="org-123", limit=10, offset=0)
126131

127-
@patch.object(DatasetClient, "get")
128-
def test_search_with_sorting(self, mock_get: MagicMock) -> None:
132+
@patch.object(DatasetClient, "_make_request")
133+
def test_search_with_sorting(self, mock_request: MagicMock) -> None:
129134
"""Test dataset search with sorting."""
130-
mock_get.return_value = {"total": 3, "results": []}
135+
mock_request.return_value = {"total": 3, "results": []}
131136

132137
result = self.client.search(query="test", sort="recent", page=1, page_size=10)
133138

134139
self.assertEqual(result["total"], 3)
135-
mock_get.assert_called_once()
140+
mock_request.assert_called_once()
136141

137142

138143
if __name__ == "__main__":

0 commit comments

Comments
 (0)