Skip to content

Commit 0b9edf2

Browse files
areeb1501areeb1501
committed
Added test cases
1 parent 417a954 commit 0b9edf2

File tree

1 file changed

+238
-0
lines changed

1 file changed

+238
-0
lines changed

‎tests/api_resources/knowledge_bases/test_indexing_jobs.py‎

Lines changed: 238 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,11 @@
55
import os
66
from typing import Any, cast
77

8+
import httpx
89
import pytest
910

1011
from gradient import Gradient, AsyncGradient
12+
from gradient import IndexingJobError, IndexingJobTimeoutError
1113
from tests.utils import assert_matches_type
1214
from gradient.types.knowledge_bases import (
1315
IndexingJobListResponse,
@@ -232,6 +234,124 @@ def test_path_params_update_cancel(self, client: Gradient) -> None:
232234
path_uuid="",
233235
)
234236

237+
@parametrize
238+
def test_wait_for_completion_raises_indexing_job_error_on_failed(self, client: Gradient, respx_mock: Any) -> None:
239+
"""Test that IndexingJobError is raised when job phase is FAILED"""
240+
job_uuid = "test-job-uuid"
241+
respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock(
242+
return_value=httpx.Response(
243+
200,
244+
json={
245+
"job": {
246+
"uuid": job_uuid,
247+
"phase": "BATCH_JOB_PHASE_FAILED",
248+
"total_items_indexed": 10,
249+
"total_items_failed": 5,
250+
}
251+
},
252+
)
253+
)
254+
255+
with pytest.raises(IndexingJobError) as exc_info:
256+
client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid)
257+
258+
assert exc_info.value.uuid == job_uuid
259+
assert exc_info.value.phase == "BATCH_JOB_PHASE_FAILED"
260+
assert "failed" in str(exc_info.value).lower()
261+
262+
@parametrize
263+
def test_wait_for_completion_raises_indexing_job_error_on_error(self, client: Gradient, respx_mock: Any) -> None:
264+
"""Test that IndexingJobError is raised when job phase is ERROR"""
265+
job_uuid = "test-job-uuid"
266+
respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock(
267+
return_value=httpx.Response(
268+
200,
269+
json={
270+
"job": {
271+
"uuid": job_uuid,
272+
"phase": "BATCH_JOB_PHASE_ERROR",
273+
}
274+
},
275+
)
276+
)
277+
278+
with pytest.raises(IndexingJobError) as exc_info:
279+
client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid)
280+
281+
assert exc_info.value.uuid == job_uuid
282+
assert exc_info.value.phase == "BATCH_JOB_PHASE_ERROR"
283+
assert "error" in str(exc_info.value).lower()
284+
285+
@parametrize
286+
def test_wait_for_completion_raises_indexing_job_error_on_cancelled(self, client: Gradient, respx_mock: Any) -> None:
287+
"""Test that IndexingJobError is raised when job phase is CANCELLED"""
288+
job_uuid = "test-job-uuid"
289+
respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock(
290+
return_value=httpx.Response(
291+
200,
292+
json={
293+
"job": {
294+
"uuid": job_uuid,
295+
"phase": "BATCH_JOB_PHASE_CANCELLED",
296+
}
297+
},
298+
)
299+
)
300+
301+
with pytest.raises(IndexingJobError) as exc_info:
302+
client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid)
303+
304+
assert exc_info.value.uuid == job_uuid
305+
assert exc_info.value.phase == "BATCH_JOB_PHASE_CANCELLED"
306+
assert "cancelled" in str(exc_info.value).lower()
307+
308+
@parametrize
309+
def test_wait_for_completion_raises_timeout_error(self, client: Gradient, respx_mock: Any) -> None:
310+
"""Test that IndexingJobTimeoutError is raised on timeout"""
311+
job_uuid = "test-job-uuid"
312+
respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock(
313+
return_value=httpx.Response(
314+
200,
315+
json={
316+
"job": {
317+
"uuid": job_uuid,
318+
"phase": "BATCH_JOB_PHASE_RUNNING",
319+
}
320+
},
321+
)
322+
)
323+
324+
with pytest.raises(IndexingJobTimeoutError) as exc_info:
325+
client.knowledge_bases.indexing_jobs.wait_for_completion(
326+
job_uuid, poll_interval=0.1, timeout=0.2
327+
)
328+
329+
assert exc_info.value.uuid == job_uuid
330+
assert exc_info.value.phase == "BATCH_JOB_PHASE_RUNNING"
331+
assert exc_info.value.timeout == 0.2
332+
333+
@parametrize
334+
def test_wait_for_completion_succeeds(self, client: Gradient, respx_mock: Any) -> None:
335+
"""Test that wait_for_completion returns successfully when job succeeds"""
336+
job_uuid = "test-job-uuid"
337+
respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock(
338+
return_value=httpx.Response(
339+
200,
340+
json={
341+
"job": {
342+
"uuid": job_uuid,
343+
"phase": "BATCH_JOB_PHASE_SUCCEEDED",
344+
"total_items_indexed": 100,
345+
"total_items_failed": 0,
346+
}
347+
},
348+
)
349+
)
350+
351+
result = client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid)
352+
assert_matches_type(IndexingJobRetrieveResponse, result, path=["response"])
353+
assert result.job.phase == "BATCH_JOB_PHASE_SUCCEEDED"
354+
235355

236356
class TestAsyncIndexingJobs:
237357
parametrize = pytest.mark.parametrize(
@@ -446,3 +566,121 @@ async def test_path_params_update_cancel(self, async_client: AsyncGradient) -> N
446566
await async_client.knowledge_bases.indexing_jobs.with_raw_response.update_cancel(
447567
path_uuid="",
448568
)
569+
570+
@parametrize
571+
async def test_wait_for_completion_raises_indexing_job_error_on_failed(self, async_client: AsyncGradient, respx_mock: Any) -> None:
572+
"""Test that IndexingJobError is raised when job phase is FAILED"""
573+
job_uuid = "test-job-uuid"
574+
respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock(
575+
return_value=httpx.Response(
576+
200,
577+
json={
578+
"job": {
579+
"uuid": job_uuid,
580+
"phase": "BATCH_JOB_PHASE_FAILED",
581+
"total_items_indexed": 10,
582+
"total_items_failed": 5,
583+
}
584+
},
585+
)
586+
)
587+
588+
with pytest.raises(IndexingJobError) as exc_info:
589+
await async_client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid)
590+
591+
assert exc_info.value.uuid == job_uuid
592+
assert exc_info.value.phase == "BATCH_JOB_PHASE_FAILED"
593+
assert "failed" in str(exc_info.value).lower()
594+
595+
@parametrize
596+
async def test_wait_for_completion_raises_indexing_job_error_on_error(self, async_client: AsyncGradient, respx_mock: Any) -> None:
597+
"""Test that IndexingJobError is raised when job phase is ERROR"""
598+
job_uuid = "test-job-uuid"
599+
respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock(
600+
return_value=httpx.Response(
601+
200,
602+
json={
603+
"job": {
604+
"uuid": job_uuid,
605+
"phase": "BATCH_JOB_PHASE_ERROR",
606+
}
607+
},
608+
)
609+
)
610+
611+
with pytest.raises(IndexingJobError) as exc_info:
612+
await async_client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid)
613+
614+
assert exc_info.value.uuid == job_uuid
615+
assert exc_info.value.phase == "BATCH_JOB_PHASE_ERROR"
616+
assert "error" in str(exc_info.value).lower()
617+
618+
@parametrize
619+
async def test_wait_for_completion_raises_indexing_job_error_on_cancelled(self, async_client: AsyncGradient, respx_mock: Any) -> None:
620+
"""Test that IndexingJobError is raised when job phase is CANCELLED"""
621+
job_uuid = "test-job-uuid"
622+
respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock(
623+
return_value=httpx.Response(
624+
200,
625+
json={
626+
"job": {
627+
"uuid": job_uuid,
628+
"phase": "BATCH_JOB_PHASE_CANCELLED",
629+
}
630+
},
631+
)
632+
)
633+
634+
with pytest.raises(IndexingJobError) as exc_info:
635+
await async_client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid)
636+
637+
assert exc_info.value.uuid == job_uuid
638+
assert exc_info.value.phase == "BATCH_JOB_PHASE_CANCELLED"
639+
assert "cancelled" in str(exc_info.value).lower()
640+
641+
@parametrize
642+
async def test_wait_for_completion_raises_timeout_error(self, async_client: AsyncGradient, respx_mock: Any) -> None:
643+
"""Test that IndexingJobTimeoutError is raised on timeout"""
644+
job_uuid = "test-job-uuid"
645+
respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock(
646+
return_value=httpx.Response(
647+
200,
648+
json={
649+
"job": {
650+
"uuid": job_uuid,
651+
"phase": "BATCH_JOB_PHASE_RUNNING",
652+
}
653+
},
654+
)
655+
)
656+
657+
with pytest.raises(IndexingJobTimeoutError) as exc_info:
658+
await async_client.knowledge_bases.indexing_jobs.wait_for_completion(
659+
job_uuid, poll_interval=0.1, timeout=0.2
660+
)
661+
662+
assert exc_info.value.uuid == job_uuid
663+
assert exc_info.value.phase == "BATCH_JOB_PHASE_RUNNING"
664+
assert exc_info.value.timeout == 0.2
665+
666+
@parametrize
667+
async def test_wait_for_completion_succeeds(self, async_client: AsyncGradient, respx_mock: Any) -> None:
668+
"""Test that wait_for_completion returns successfully when job succeeds"""
669+
job_uuid = "test-job-uuid"
670+
respx_mock.get(f"{base_url}/v2/gen-ai/indexing_jobs/{job_uuid}").mock(
671+
return_value=httpx.Response(
672+
200,
673+
json={
674+
"job": {
675+
"uuid": job_uuid,
676+
"phase": "BATCH_JOB_PHASE_SUCCEEDED",
677+
"total_items_indexed": 100,
678+
"total_items_failed": 0,
679+
}
680+
},
681+
)
682+
)
683+
684+
result = await async_client.knowledge_bases.indexing_jobs.wait_for_completion(job_uuid)
685+
assert_matches_type(IndexingJobRetrieveResponse, result, path=["response"])
686+
assert result.job.phase == "BATCH_JOB_PHASE_SUCCEEDED"

0 commit comments

Comments
 (0)