def test_add_item_happy_path(self, mock_client, mock_query): current_session_key = datastore.key.Key("session", "current", project=0) current_session_vars = { "region": "us_va", "scrape_type": constants.ScrapeType.SNAPSHOT, "phase": scrape_phase.ScrapePhase.START, "start": fix_dt(datetime(2014, 8, 31)), } current_session = ScrapeSession.new(current_session_key, **current_session_vars) prior_session = ScrapeSession.new( datastore.key.Key("session", "prior", project=0), region="us_ny", scrape_type=constants.ScrapeType.SNAPSHOT, start=fix_dt(datetime(2014, 8, 17)), phase=scrape_phase.ScrapePhase.SCRAPE, ) wire_sessions_to_query( mock_client, mock_query, [current_session, prior_session] ) assert sessions.add_docket_item_to_current_session( "alpha", ScrapeKey("us_va", constants.ScrapeType.SNAPSHOT) ) current_session_vars.update({"docket_ack_id": "alpha"}) expected_session = ScrapeSession.new( current_session_key, **current_session_vars ) mock_client.return_value.put.assert_called_with(expected_session.to_entity())
def test_add_item_happy_path(self, mock_client, mock_query): current_session_key = datastore.key.Key('session', 'current', project=0) current_session_vars = { 'region': 'us_va', 'scrape_type': constants.ScrapeType.SNAPSHOT, 'phase': scrape_phase.ScrapePhase.START, 'start': fix_dt(datetime(2014, 8, 31)) } current_session = ScrapeSession.new(current_session_key, **current_session_vars) prior_session = ScrapeSession.new( datastore.key.Key('session', 'prior', project=0), region='us_ny', scrape_type=constants.ScrapeType.SNAPSHOT, start=fix_dt( datetime(2014, 8, 17)), phase=scrape_phase.ScrapePhase.SCRAPE, ) wire_sessions_to_query( mock_client, mock_query, [current_session, prior_session]) assert sessions.add_docket_item_to_current_session( "alpha", ScrapeKey("us_va", constants.ScrapeType.SNAPSHOT)) current_session_vars.update({'docket_ack_id': 'alpha'}) expected_session = ScrapeSession.new( current_session_key, **current_session_vars ) mock_client.return_value.put.assert_called_with( expected_session.to_entity())
def iterate_docket_item(scrape_key, return_immediately=False): """Leases new docket item, updates current session, returns item contents Pulls an arbitrary new item from the docket type provided, adds it to the current session info, and returns the payload of the docket item. This payload should be an entity fit to scrape, or information suitable for retrieving an entity fit to scrape, depending on scrape type. Args: scrape_key: (ScrapeKey) The scraper to retrieve a docket item for return_immediately: (bool) Whether to return immediately or to wait for a bounded period of time for a message to enter the docket. Returns: The payload of the next docket item, if successfully retrieved and added to the current session for this region and scrape type. If not retrieved or not successfully added to the session, returns None. """ docket_item = docket.get_new_docket_item( scrape_key, return_immediately=return_immediately ) if not docket_item: logging.info("No items in docket for [%s]. Ending scrape.", scrape_key) return None item_content = json.loads(docket_item.message.data.decode()) item_added = sessions.add_docket_item_to_current_session( docket_item.ack_id, scrape_key ) if not item_added: logging.error( "Failed to update session for scraper [%s] " "with docket item [%s].", scrape_key, str(item_content), ) return None return item_content
def test_add_item_no_open_sessions(self, _mock_client): assert not sessions.add_docket_item_to_current_session( "alpha", ScrapeKey("us_va", constants.ScrapeType.SNAPSHOT) )