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

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

    counts = export_g8_suppliers.service_counts(record)

    assert counts == [
        ('completed_saas', 2),
        ('completed_paas', 1),
        ('completed_iaas', 0),
        ('completed_scs', 0),
        ('draft_saas', 0),
        ('draft_paas', 3),
        ('draft_iaas', 0),
        ('draft_scs', 2),
    ]
def test_add_draft_counts(mock_data_client):
    mock_data_client.find_draft_services_iter.return_value = [
        {'status': 'submitted', 'lot': 'saas'},
        {'status': 'submitted', 'lot': 'saas'},
        {'status': 'submitted', 'lot': 'paas'},
        {'status': 'not-submitted', 'lot': 'saas'},
        {'status': 'not-submitted', 'lot': 'paas'},
        {'status': 'not-submitted', 'lot': 'paas'},
        {'status': 'not-submitted', 'lot': 'paas'},
        {'status': 'published', 'lot': 'scs'},  # anything not submitted or failed is considered draft
    ]

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

    record = draft_counts_adder({'supplier': {'id': 1}})
    print("COUNTS: {}".format(record['counts']))
    assert record['counts'] == {
        'completed': {
            'saas': 2, 'paas': 1,
            'iaas': 0, 'scs': 0,
        },
        'draft': {
            'saas': 1, 'paas': 3,
            'iaas': 0, 'scs': 1,
        },
    }