def test_count_and_set_headers_flex(): """Verify that we are setting the correct flex headers""" reader = csvReader.CsvReader() csv_schema = [Mock(name_short='some_col'), Mock(name_short='other')] header_row = ['ignored', 'flex_my_col', 'some_col', 'flex_other', 'some_col'] result = reader.count_and_set_headers(csv_schema, header_row) assert result == {'some_col': 2, 'other': 0} assert reader.expected_headers == [None, None, 'some_col', None, 'some_col'] assert reader.flex_headers == [None, 'flex_my_col', None, 'flex_other', None]
def test_get_next_record_flex(): """Verify that we get a list of FlexFields if present""" reader = csvReader.CsvReader() reader.column_count = 6 reader.expected_headers = ['a', 'b', 'c', None, None, None] reader.flex_headers = [None, None, None, 'flex_d', 'flex_e', None] reader.csv_reader = csv.reader(['A,"B\n",C,D,E,F'], dialect='excel', delimiter=',') return_dict, flex_fields = reader.get_next_record() assert return_dict == {'a': 'A', 'b': 'B\n', 'c': 'C'} assert len(flex_fields) == 2 assert flex_fields[0].header == 'flex_d' assert flex_fields[0].cell == 'D' assert flex_fields[1].header == 'flex_e' assert flex_fields[1].cell == 'E'