From e7c8c5bcf9d0b3d529e086489850025107f57121 Mon Sep 17 00:00:00 2001
From: Harrison Deng <yunyangdeng@outlook.com>
Date: Fri, 17 Jan 2025 16:56:52 +0000
Subject: [PATCH] Changed allele CSV output to non-lists

---
 src/automlst/engine/data/local/csv.py        | 12 ++++++------
 tests/automlst/engine/data/local/test_csv.py | 12 ++++++++++++
 2 files changed, 18 insertions(+), 6 deletions(-)
 create mode 100644 tests/automlst/engine/data/local/test_csv.py

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