def test_find_suppliers_filtered_by_supplier_ids(mock_data_client):
    mock_data_client.get_interested_suppliers.return_value = {
        'interestedSuppliers': [4, 3, 2]
    }

    records = list(export_dos_suppliers.find_suppliers(mock_data_client, 'framework-slug', supplier_ids=[2, 4]))

    assert records == [
        {'supplier_id': 4}, {'supplier_id': 2},
    ]
def find_all_labs(client):
    pool = ThreadPool(20)
    records = find_suppliers(client, FRAMEWORK_SLUG)
    records = pool.imap(add_framework_info(client, FRAMEWORK_SLUG), records)
    records = filter(lambda record: record['onFramework'], records)
    records = pool.imap(add_draft_services(client, FRAMEWORK_SLUG), records)
    services = itertools.chain.from_iterable(record['services'] for record in records)
    services = filter(
        lambda record: record['lot'] == 'user-research-studios' and record['status'] == 'submitted',
        services)

    return services
def test_find_suppliers_produces_results_with_supplier_ids(mock_data_client):
    mock_data_client.get_interested_suppliers.return_value = {
        'interestedSuppliers': [4, 3, 2]
    }

    records = list(export_dos_suppliers.find_suppliers(mock_data_client, 'framework-slug'))

    mock_data_client.get_interested_suppliers.assert_has_calls([
        call('framework-slug')
    ])
    assert records == [
        {'supplier_id': 4}, {'supplier_id': 3}, {'supplier_id': 2}
    ]