|
9 | 9 | """ |
10 | 10 |
|
11 | 11 | import time |
| 12 | +from datetime import datetime |
| 13 | +from unittest.mock import MagicMock, patch |
12 | 14 |
|
13 | 15 | import pytest |
14 | 16 | from groundlight import ExperimentalApi |
@@ -51,6 +53,126 @@ def test_list_priming_groups_returns_list(gl_experimental: ExperimentalApi): |
51 | 53 | assert all(isinstance(g, PrimingGroup) for g in groups) |
52 | 54 |
|
53 | 55 |
|
| 56 | +# --------------------------------------------------------------------------- |
| 57 | +# create / get / delete (mocked) |
| 58 | +# --------------------------------------------------------------------------- |
| 59 | + |
| 60 | +_MOCK_PIPELINE = { |
| 61 | + "id": "pipe_mock0000000000000001", |
| 62 | + "pipeline_config": "never-review", |
| 63 | + "is_active_pipeline": True, |
| 64 | + "is_edge_pipeline": False, |
| 65 | + "is_unclear_pipeline": False, |
| 66 | + "is_oodd_pipeline": False, |
| 67 | + "is_enabled": True, |
| 68 | + "created_at": "2024-01-01T00:00:00Z", |
| 69 | + "trained_at": "2024-01-01T01:00:00Z", |
| 70 | +} |
| 71 | + |
| 72 | +_MOCK_PG = { |
| 73 | + "id": "pgp_mock0000000000000001", |
| 74 | + "name": "test-primer", |
| 75 | + "is_global": False, |
| 76 | + "canonical_query": "Is there a dog?", |
| 77 | + "active_pipeline_config": None, |
| 78 | + "priming_group_specific_shadow_pipeline_configs": None, |
| 79 | + "disable_shadow_pipelines": False, |
| 80 | + "created_at": "2024-01-01T00:00:00Z", |
| 81 | +} |
| 82 | + |
| 83 | + |
| 84 | +def _mock_response(status_code: int, json_body=None): |
| 85 | + resp = MagicMock() |
| 86 | + resp.status_code = status_code |
| 87 | + resp.json.return_value = json_body or {} |
| 88 | + resp.raise_for_status = MagicMock() |
| 89 | + if status_code >= 400: |
| 90 | + resp.raise_for_status.side_effect = Exception(f"HTTP {status_code}") |
| 91 | + return resp |
| 92 | + |
| 93 | + |
| 94 | +def test_create_priming_group_mocked(gl_experimental: ExperimentalApi): |
| 95 | + with patch("groundlight.experimental_api.requests.post") as mock_post: |
| 96 | + mock_post.return_value = _mock_response(201, _MOCK_PG) |
| 97 | + |
| 98 | + pg = gl_experimental.create_priming_group( |
| 99 | + name="test-primer", |
| 100 | + source_ml_pipeline_id="pipe_mock0000000000000001", |
| 101 | + canonical_query="Is there a dog?", |
| 102 | + ) |
| 103 | + |
| 104 | + assert isinstance(pg, PrimingGroup) |
| 105 | + assert pg.id == "pgp_mock0000000000000001" |
| 106 | + assert pg.name == "test-primer" |
| 107 | + assert pg.canonical_query == "Is there a dog?" |
| 108 | + assert pg.is_global is False |
| 109 | + mock_post.assert_called_once() |
| 110 | + _, kwargs = mock_post.call_args |
| 111 | + assert kwargs["json"]["name"] == "test-primer" |
| 112 | + assert kwargs["json"]["source_ml_pipeline_id"] == "pipe_mock0000000000000001" |
| 113 | + |
| 114 | + |
| 115 | +def test_create_priming_group_disable_shadow_pipelines_mocked(gl_experimental: ExperimentalApi): |
| 116 | + pg_data = {**_MOCK_PG, "name": "test-primer-noshadow", "disable_shadow_pipelines": True} |
| 117 | + with patch("groundlight.experimental_api.requests.post") as mock_post: |
| 118 | + mock_post.return_value = _mock_response(201, pg_data) |
| 119 | + |
| 120 | + pg = gl_experimental.create_priming_group( |
| 121 | + name="test-primer-noshadow", |
| 122 | + source_ml_pipeline_id="pipe_mock0000000000000001", |
| 123 | + disable_shadow_pipelines=True, |
| 124 | + ) |
| 125 | + |
| 126 | + assert pg.disable_shadow_pipelines is True |
| 127 | + _, kwargs = mock_post.call_args |
| 128 | + assert kwargs["json"]["disable_shadow_pipelines"] is True |
| 129 | + |
| 130 | + |
| 131 | +def test_get_priming_group_mocked(gl_experimental: ExperimentalApi): |
| 132 | + with patch("groundlight.experimental_api.requests.get") as mock_get: |
| 133 | + mock_get.return_value = _mock_response(200, _MOCK_PG) |
| 134 | + |
| 135 | + pg = gl_experimental.get_priming_group("pgp_mock0000000000000001") |
| 136 | + |
| 137 | + assert isinstance(pg, PrimingGroup) |
| 138 | + assert pg.id == "pgp_mock0000000000000001" |
| 139 | + assert pg.name == "test-primer" |
| 140 | + |
| 141 | + |
| 142 | +def test_get_priming_group_unknown_raises_mocked(gl_experimental: ExperimentalApi): |
| 143 | + with patch("groundlight.experimental_api.requests.get") as mock_get: |
| 144 | + mock_get.return_value = _mock_response(404) |
| 145 | + |
| 146 | + with pytest.raises(NotFoundError): |
| 147 | + gl_experimental.get_priming_group("pgp_doesnotexist000000000000") |
| 148 | + |
| 149 | + |
| 150 | +def test_delete_priming_group_mocked(gl_experimental: ExperimentalApi): |
| 151 | + with patch("groundlight.experimental_api.requests.get") as mock_get, patch( |
| 152 | + "groundlight.experimental_api.requests.delete" |
| 153 | + ) as mock_delete: |
| 154 | + mock_get.return_value = _mock_response(404) |
| 155 | + mock_delete.return_value = _mock_response(204) |
| 156 | + |
| 157 | + gl_experimental.delete_priming_group("pgp_mock0000000000000001") |
| 158 | + |
| 159 | + # After deletion, getting it should raise NotFoundError |
| 160 | + with pytest.raises(NotFoundError): |
| 161 | + gl_experimental.get_priming_group("pgp_mock0000000000000001") |
| 162 | + |
| 163 | + mock_delete.assert_called_once() |
| 164 | + |
| 165 | + |
| 166 | +def test_created_priming_group_appears_in_list_mocked(gl_experimental: ExperimentalApi): |
| 167 | + list_response = {"results": [_MOCK_PG]} |
| 168 | + with patch("groundlight.experimental_api.requests.get") as mock_get: |
| 169 | + mock_get.return_value = _mock_response(200, list_response) |
| 170 | + |
| 171 | + groups = gl_experimental.list_priming_groups() |
| 172 | + |
| 173 | + assert any(g.id == "pgp_mock0000000000000001" for g in groups) |
| 174 | + |
| 175 | + |
54 | 176 | # --------------------------------------------------------------------------- |
55 | 177 | # create / get / delete (expensive — require a trained pipeline) |
56 | 178 | # --------------------------------------------------------------------------- |
|
0 commit comments