def test_service_counts(mock_data_client):
    # Count services
    mock_data_client.find_draft_services_iter.return_value = [
        {'status': 'submitted', 'lot': 'digital-outcomes'},
        {'status': 'submitted', 'lot': 'digital-outcomes'},
        {'status': 'submitted', 'lot': 'digital-specialists'},
        {'status': 'failed', 'lot': 'digital-outcomes'},
        {'status': 'failed', 'lot': 'digital-outcomes'},
        {'status': 'not-submitted', 'lot': 'digital-specialists'},
        {'status': 'not-submitted', 'lot': 'digital-specialists'},
        {'status': 'published', 'lot': 'digital-specialists'},  # anything not submitted or failed is considered draft
    ]
    draft_counts_adder = export_dos_suppliers.add_draft_counts(mock_data_client, 'framework-slug')

    record = draft_counts_adder({'supplier': {'id': 1}})

    counts = export_dos_suppliers.service_counts(record)

    assert counts == [
        ('completed_digital-outcomes', 2),
        ('completed_digital-specialists', 1),
        ('completed_user-research-studios', 0),
        ('completed_user-research-participants', 0),
        ('failed_digital-outcomes', 2),
        ('failed_digital-specialists', 0),
        ('failed_user-research-studios', 0),
        ('failed_user-research-participants', 0),
        ('draft_digital-outcomes', 0),
        ('draft_digital-specialists', 3),
        ('draft_user-research-studios', 0),
        ('draft_user-research-participants', 0),
    ]
def test_add_draft_counts(mock_data_client):
    mock_data_client.find_draft_services_iter.return_value = [
        {'status': 'submitted', 'lot': 'digital-outcomes'},
        {'status': 'submitted', 'lot': 'digital-outcomes'},
        {'status': 'submitted', 'lot': 'digital-specialists'},
        {'status': 'failed', 'lot': 'digital-outcomes'},
        {'status': 'failed', 'lot': 'digital-outcomes'},
        {'status': 'not-submitted', 'lot': 'digital-specialists'},
        {'status': 'not-submitted', 'lot': 'digital-specialists'},
        {'status': 'published', 'lot': 'digital-specialists'},  # anything not submitted or failed is considered draft
    ]

    draft_counts_adder = export_dos_suppliers.add_draft_counts(mock_data_client, 'framework-slug')

    record = draft_counts_adder({'supplier': {'id': 1}})

    assert record['counts'] == {
        'completed': {
            'digital-outcomes': 2, 'digital-specialists': 1,
            'user-research-studios': 0, 'user-research-participants': 0,
        },
        'failed': {
            'digital-outcomes': 2, 'digital-specialists': 0,
            'user-research-studios': 0, 'user-research-participants': 0,
        },
        'draft': {
            'digital-outcomes': 0, 'digital-specialists': 3,
            'user-research-studios': 0, 'user-research-participants': 0,
        },
    }