<commit_before>import os
import pandas as pd

from dataactbroker.scripts import update_historical_duns
from dataactcore.config import CONFIG_BROKER
from dataactcore.utils.duns import DUNS_COLUMNS, EXCLUDE_FROM_API
from dataactcore.models.domainModels import DUNS, HistoricDUNS


def test_remove_existing_duns(database):
    """ Testing the removing existing duns function"""
    sess = database.session
    # of the duns 000000001-000000009, half of them are in the database
    all_duns = ['00000000{}'.format(x) for x in range(0, 10)]
    existing_duns = all_duns[: 4]
    data = pd.DataFrame.from_dict({'awardee_or_recipient_uniqu': all_duns})
    for duns in existing_duns:
        sess.add(DUNS(awardee_or_recipient_uniqu=duns))
    sess.commit()

    # confirm that the dataframe returned only has half the duns
    expected_duns = list(set(existing_duns) ^ set(all_duns))
    new_df = update_historical_duns.remove_existing_duns(data, sess)
    assert sorted(expected_duns) == sorted(new_df['awardee_or_recipient_uniqu'].tolist())


def mock_get_duns_props_from_sam(duns_list):
    """ Mock function for get_duns_props as we can't connect to the SAM service """
    request_cols = [col for col in DUNS_COLUMNS if col not in EXCLUDE_FROM_API]
    columns = request_cols
    results = pd.DataFrame(columns=columns)

    duns_mappings = {
        '000000001': {
            'awardee_or_recipient_uniqu': '000000001',
            'uei': 'A1',
            'legal_business_name': 'Legal Name 1',
            'dba_name': 'Name 1',
            'entity_structure': '1A',
            'ultimate_parent_unique_ide': '999999999',
            'ultimate_parent_uei': 'Z9',
            'ultimate_parent_legal_enti': 'Parent Legal Name 1',
            'address_line_1': 'Test address 1',
            'address_line_2': 'Test address 2',
            'city': 'Test city',
            'state': 'Test state',
            'zip': 'Test zip',
            'zip4': 'Test zip4',
            'country_code': 'Test country',
            'congressional_district': 'Test congressional district',
            'business_types_codes': [['A', 'B', 'C']],<commit_msg>Accounting for API arg for get_duns_props_from sam<commit_after>import os
import pandas as pd

from dataactbroker.scripts import update_historical_duns
from dataactcore.config import CONFIG_BROKER
from dataactcore.utils.duns import DUNS_COLUMNS, EXCLUDE_FROM_API
from dataactcore.models.domainModels import DUNS, HistoricDUNS


def test_remove_existing_duns(database):
    """ Testing the removing existing duns function"""
    sess = database.session
    # of the duns 000000001-000000009, half of them are in the database
    all_duns = ['00000000{}'.format(x) for x in range(0, 10)]
    existing_duns = all_duns[: 4]
    data = pd.DataFrame.from_dict({'awardee_or_recipient_uniqu': all_duns})
    for duns in existing_duns:
        sess.add(DUNS(awardee_or_recipient_uniqu=duns))
    sess.commit()

    # confirm that the dataframe returned only has half the duns
    expected_duns = list(set(existing_duns) ^ set(all_duns))
    new_df = update_historical_duns.remove_existing_duns(data, sess)
    assert sorted(expected_duns) == sorted(new_df['awardee_or_recipient_uniqu'].tolist())


def mock_get_duns_props_from_sam(duns_list, api='entity'):
    """ Mock function for get_duns_props as we can't connect to the SAM service """
    request_cols = [col for col in DUNS_COLUMNS if col not in EXCLUDE_FROM_API]
    columns = request_cols
    results = pd.DataFrame(columns=columns)

    duns_mappings = {
        '000000001': {
            'awardee_or_recipient_uniqu': '000000001',
            'uei': 'A1',
            'legal_business_name': 'Legal Name 1',
            'dba_name': 'Name 1',
            'entity_structure': '1A',
            'ultimate_parent_unique_ide': '999999999',
            'ultimate_parent_uei': 'Z9',
            'ultimate_parent_legal_enti': 'Parent Legal Name 1',
            'address_line_1': 'Test address 1',
            'address_line_2': 'Test address 2',
            'city': 'Test city',
            'state': 'Test state',
            'zip': 'Test zip',
            'zip4': 'Test zip4',
            'country_code': 'Test country',
            'congressional_district': 'Test congressional district',
            'business_types_codes': [['A', 'B', 'C']],
{'fraction_non_alphanumeric': 0.1054335260115607, 'fraction_numerical': 0.03190751445086705, 'mean_word_length': 3.326, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 12, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1689274', 'n_tokens_mistral': 1533, 'n_tokens_neox': 1386, 'n_words': 340}