|
5 | 5 | import pytest |
6 | 6 | from aiohttp import ClientSession |
7 | 7 |
|
| 8 | +from tests.helpers.test_batch_creation_parameters import TestBatchCreationParameters, TestURLCreationParameters |
8 | 9 | from source_collectors.muckrock.MuckrockAPIInterface import MuckrockAPIInterface, AgencyLookupResponseType, AgencyLookupResponse |
9 | 10 | from collector_db.models import Agency, AutomatedUrlAgencySuggestion |
10 | | -from collector_manager.enums import CollectorType |
| 11 | +from collector_manager.enums import CollectorType, URLStatus |
11 | 12 | from core.DTOs.TaskOperatorRunInfo import TaskOperatorOutcome |
12 | 13 | from core.DTOs.URLAgencySuggestionInfo import URLAgencySuggestionInfo |
13 | 14 | from core.classes.task_operators.AgencyIdentificationTaskOperator import AgencyIdentificationTaskOperator |
|
20 | 21 | from pdap_api_client.DTOs import MatchAgencyResponse, MatchAgencyInfo |
21 | 22 | from pdap_api_client.PDAPClient import PDAPClient |
22 | 23 | from pdap_api_client.enums import MatchAgencyResponseStatus |
23 | | -from tests.helpers.DBDataCreator import DBDataCreator, BatchURLCreationInfo |
| 24 | +from tests.helpers.DBDataCreator import DBDataCreator, BatchURLCreationInfo, BatchURLCreationInfoV2 |
24 | 25 |
|
25 | 26 | sample_agency_suggestions = [ |
26 | 27 | URLAgencySuggestionInfo( |
@@ -103,8 +104,25 @@ async def mock_run_subtask( |
103 | 104 | CollectorType.MUCKROCK_ALL_SEARCH, |
104 | 105 | CollectorType.CKAN |
105 | 106 | ]: |
106 | | - creation_info: BatchURLCreationInfo = await db_data_creator.batch_and_urls(strategy=strategy, url_count=1, with_html_content=True) |
107 | | - d[strategy] = creation_info.url_ids[0] |
| 107 | + # Create two URLs for each, one pending and one errored |
| 108 | + creation_info: BatchURLCreationInfoV2 = await db_data_creator.batch_v2( |
| 109 | + parameters=TestBatchCreationParameters( |
| 110 | + strategy=strategy, |
| 111 | + urls=[ |
| 112 | + TestURLCreationParameters( |
| 113 | + count=1, |
| 114 | + status=URLStatus.PENDING, |
| 115 | + with_html_content=True |
| 116 | + ), |
| 117 | + TestURLCreationParameters( |
| 118 | + count=1, |
| 119 | + status=URLStatus.ERROR, |
| 120 | + with_html_content=True |
| 121 | + ) |
| 122 | + ] |
| 123 | + ) |
| 124 | + ) |
| 125 | + d[strategy] = creation_info.url_creation_infos[URLStatus.PENDING].url_mappings[0].url_id |
108 | 126 |
|
109 | 127 |
|
110 | 128 | # Confirm meets prerequisites |
|
0 commit comments