diff --git a/src/automlst/engine/data/local/csv.py b/src/automlst/engine/data/local/csv.py index 8bef3d0..832ee47 100644 --- a/src/automlst/engine/data/local/csv.py +++ b/src/automlst/engine/data/local/csv.py @@ -1,7 +1,6 @@ import csv -from io import TextIOWrapper from os import PathLike -from typing import AsyncIterable, Iterable, Mapping, Sequence, Union +from typing import AsyncIterable, Mapping, Sequence, Union from automlst.engine.data.structures.mlst import Allele, MLSTProfile @@ -11,10 +10,11 @@ def dict_loci_alleles_variants_from_loci(alleles_map: Mapping[str, Sequence[Alle for loci, alleles in alleles_map.items(): if len(alleles) == 1: result_dict[loci] = alleles[0].allele_variant - for allele in alleles: - result_locis = list() - result_locis.append(allele.allele_variant) - result_dict[loci] = result_locis + else: + for allele in alleles: + result_locis = list() + result_locis.append(allele.allele_variant) + result_dict[loci] = result_locis return result_dict diff --git a/tests/automlst/engine/data/local/test_csv.py b/tests/automlst/engine/data/local/test_csv.py new file mode 100644 index 0000000..2095990 --- /dev/null +++ b/tests/automlst/engine/data/local/test_csv.py @@ -0,0 +1,12 @@ +from automlst.engine.data.local.csv import dict_loci_alleles_variants_from_loci +from automlst.engine.data.structures.mlst import Allele + + +def test_dict_loci_alleles_variants_from_loci_single_loci_not_list(): + alleles_map = { + "adk": [Allele("adk", "1", None)] + } + results = dict_loci_alleles_variants_from_loci(alleles_map) + for loci, variant in results.items(): + assert isinstance(variant, str) + assert variant == "1" \ No newline at end of file