diff --git a/src/automlst/engine/data/local/csv.py b/src/automlst/engine/data/local/csv.py index 832ee47..d8e97f5 100644 --- a/src/automlst/engine/data/local/csv.py +++ b/src/automlst/engine/data/local/csv.py @@ -11,8 +11,8 @@ def dict_loci_alleles_variants_from_loci(alleles_map: Mapping[str, Sequence[Alle if len(alleles) == 1: result_dict[loci] = alleles[0].allele_variant else: + result_locis = list() 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 index 2095990..584523d 100644 --- a/tests/automlst/engine/data/local/test_csv.py +++ b/tests/automlst/engine/data/local/test_csv.py @@ -9,4 +9,13 @@ def test_dict_loci_alleles_variants_from_loci_single_loci_not_list(): 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 + assert variant == "1" + +def test_dict_loci_alleles_variants_from_loci_multi_loci_is_list(): + alleles_map = { + "adk": [Allele("adk", "1", None), Allele("adk", "2", None)] + } + results = dict_loci_alleles_variants_from_loci(alleles_map) + for loci, variant in results.items(): + assert isinstance(variant, list) + assert len(variant) == 2 \ No newline at end of file