|
23 | 23 | logger, |
24 | 24 | ) |
25 | 25 | from braintrust.db_fields import AUDIT_METADATA_FIELD |
| 26 | +from braintrust.git_fields import GitMetadataSettings, RepoInfo |
26 | 27 | from braintrust.gitutil import get_repo_info |
27 | 28 | from braintrust.id_gen import OTELIDGenerator, get_id_generator |
28 | 29 | from braintrust.logger import ( |
@@ -175,6 +176,67 @@ def test_init_disable_atexit_flush(self): |
175 | 176 | _HTTPBackgroundLogger(LazyValue(api_con_response, use_mutex=False)) # type: ignore |
176 | 177 | mock_register.assert_not_called() |
177 | 178 |
|
| 179 | + def test_init_without_git_metadata_override_uses_org_policy(self): |
| 180 | + for org_settings in ( |
| 181 | + GitMetadataSettings(collect="all"), |
| 182 | + GitMetadataSettings(collect="some", fields=["commit", "branch"]), |
| 183 | + ): |
| 184 | + with self.subTest(org_settings=org_settings): |
| 185 | + mock_conn = MagicMock() |
| 186 | + mock_conn.post_json.return_value = { |
| 187 | + "project": {"id": "test-project-id", "name": "test-project"}, |
| 188 | + "experiment": {"id": "test-exp-id", "name": "test-exp"}, |
| 189 | + } |
| 190 | + |
| 191 | + simulate_login() |
| 192 | + logger._state.git_metadata_settings = org_settings |
| 193 | + with patch.object(logger._state, "app_conn", return_value=mock_conn): |
| 194 | + with patch( |
| 195 | + "braintrust.logger.get_repo_info", return_value=RepoInfo(commit="abc123") |
| 196 | + ) as mock_get_repo_info: |
| 197 | + exp = braintrust.init(project="test-project") |
| 198 | + exp._lazy_metadata.get() |
| 199 | + |
| 200 | + actual_settings = mock_get_repo_info.call_args.args[0] |
| 201 | + assert actual_settings == org_settings |
| 202 | + |
| 203 | + def test_init_git_metadata_override_merges_with_org_policy(self): |
| 204 | + mock_conn = MagicMock() |
| 205 | + mock_conn.post_json.return_value = { |
| 206 | + "project": {"id": "test-project-id", "name": "test-project"}, |
| 207 | + "experiment": {"id": "test-exp-id", "name": "test-exp"}, |
| 208 | + } |
| 209 | + |
| 210 | + simulate_login() |
| 211 | + logger._state.git_metadata_settings = GitMetadataSettings(collect="some", fields=["commit", "branch"]) |
| 212 | + with patch.object(logger._state, "app_conn", return_value=mock_conn): |
| 213 | + with patch("braintrust.logger.get_repo_info", return_value=None) as mock_get_repo_info: |
| 214 | + exp = braintrust.init( |
| 215 | + project="test-project", |
| 216 | + git_metadata_settings=GitMetadataSettings(collect="some", fields=["commit"]), |
| 217 | + ) |
| 218 | + exp._lazy_metadata.get() |
| 219 | + |
| 220 | + actual_settings = mock_get_repo_info.call_args.args[0] |
| 221 | + assert actual_settings == GitMetadataSettings(collect="some", fields=["commit"]) |
| 222 | + |
| 223 | + def test_init_without_git_metadata_policy_collects_none(self): |
| 224 | + mock_conn = MagicMock() |
| 225 | + mock_conn.post_json.return_value = { |
| 226 | + "project": {"id": "test-project-id", "name": "test-project"}, |
| 227 | + "experiment": {"id": "test-exp-id", "name": "test-exp"}, |
| 228 | + } |
| 229 | + |
| 230 | + simulate_login() |
| 231 | + assert logger._state.git_metadata_settings is None |
| 232 | + with patch.object(logger._state, "app_conn", return_value=mock_conn): |
| 233 | + with patch("braintrust.logger.get_repo_info", return_value=None) as mock_get_repo_info: |
| 234 | + exp = braintrust.init(project="test-project") |
| 235 | + exp._lazy_metadata.get() |
| 236 | + |
| 237 | + actual_settings = mock_get_repo_info.call_args.args[0] |
| 238 | + assert actual_settings == GitMetadataSettings(collect="none") |
| 239 | + |
178 | 240 | def test_init_with_saved_parameters_attaches_reference(self): |
179 | 241 | mock_conn = MagicMock() |
180 | 242 | mock_conn.post_json.return_value = { |
|
0 commit comments