add methods and test to get csv from genes endpoint

backend/metagenedb/api/catalog/views/gene.py

-import hashlib
 from io import StringIO
 
-from django.core.cache import cache
 from django.conf import settings
 from django.http import HttpResponse
 from drf_yasg.utils import swagger_auto_schema
@@ -54,13 +52,51 @@ class GeneViewSet(BulkViewSet):
             response['Content-Disposition'] = 'attachment; filename=%s' % filename
             return response
 
+    def _extract_taxonomy_info(self, gene):
+        if gene.taxonomy is None:
+            return ['', '', '']
+        return [gene.taxonomy.tax_id, gene.taxonomy.name, gene.taxonomy.rank]
+
+    def _extract_function_info(self, gene):
+        if not gene.functions.all():
+            return ['', '']
+        function_ids = {
+            'kegg': [],
+            'eggnog': []
+        }
+        for function in gene.functions.all():
+            function_ids.get(function.source).append(function.function_id)
+        return [
+            ';'.join(function_ids['kegg']),
+            ';'.join(function_ids['eggnog'])
+        ]
+
+    def _get_metadata_line(self, gene):
+        """
+        Transform gene content to a line for metadata extract
+        """
+        gene_items = [
+            gene.gene_id, gene.name, gene.source, gene.length,
+        ]
+        gene_items = gene_items + self._extract_taxonomy_info(gene)
+        gene_items = gene_items + self._extract_function_info(gene)
+        return ','.join([str(item) for item in gene_items])
+
     def _build_csv_response(self):
         queryset = self.filter_queryset(self.get_queryset())
         if self._check_too_many_genes(queryset):
             return self.too_many_genes_error_response
         with StringIO() as csv_file:
+            # Write header
+            header = ",".join([
+                'gene_id', 'gene_name', 'gene_source', 'length', 'tax_id', 'tax_name', 'tax_rank',
+                'kegg_id', 'eggnog_id',
+            ])
+            csv_file.write(header)
+            for gene in queryset.iterator():
+                csv_file.write(f"{self._get_metadata_line(gene)}\n")
             # generate the file
-            response = HttpResponse("Incoming", content_type='text/csv')
+            response = HttpResponse(csv_file.getvalue(), content_type='text/csv')
             filename = 'metagenedb.csv'
             response['Content-Disposition'] = 'attachment; filename=%s' % filename
             return response

backend/metagenedb/api/catalog/views/test_gene.py

@@ -2,6 +2,20 @@ from django.test import TestCase
 from django.urls import reverse
 from rest_framework import status
 
+from metagenedb.api.catalog.views import GeneViewSet
+from metagenedb.apps.catalog.factory import (
+    GeneFactory, GeneWithEggNOGFactory, GeneWithKeggFactory, GeneWithTaxonomyFactory
+)
+
+
+class GeneViewSetMock(GeneViewSet):
+    """
+    Make mock in case complex instantiation occurs since tests here are independant from the class itself
+    """
+
+    def __init__(self):
+        pass
+
 
 class TestGenes(TestCase):
 
@@ -12,3 +26,54 @@ class TestGenes(TestCase):
         url = reverse('api:catalog:v1:genes-list')
         resp = self.client.get(url)
         self.assertEqual(resp.status_code, status.HTTP_200_OK)
+
+    def test_get_metadata_line_no_functions(self):
+        gene = GeneFactory()
+        expected_items = [
+            gene.gene_id, gene.name, gene.source, gene.length,
+            '', '', '', '', ''
+        ]
+        expected_line = ','.join([str(item) for item in expected_items])
+        # Make test with method from GeneViewSet
+        viewset = GeneViewSetMock()
+        tested_line = viewset._get_metadata_line(gene)
+        self.assertEqual(tested_line, expected_line)
+
+    def test_get_metadata_line_with_taxonomy(self):
+        gene = GeneWithTaxonomyFactory()
+        expected_items = [
+            gene.gene_id, gene.name, gene.source, gene.length,
+            gene.taxonomy.tax_id, gene.taxonomy.name, gene.taxonomy.rank,
+            '', ''
+        ]
+        expected_line = ','.join([str(item) for item in expected_items])
+        # Make test with method from GeneViewSet
+        viewset = GeneViewSetMock()
+        tested_line = viewset._get_metadata_line(gene)
+        self.assertEqual(tested_line, expected_line)
+
+    def test_get_metadata_line_with_kegg(self):
+        gene = GeneWithKeggFactory()
+        expected_items = [
+            gene.gene_id, gene.name, gene.source, gene.length,
+            '', '', '',
+            gene.functions.all()[0].function_id, ''
+        ]
+        expected_line = ','.join([str(item) for item in expected_items])
+        # Make test with method from GeneViewSet
+        viewset = GeneViewSetMock()
+        tested_line = viewset._get_metadata_line(gene)
+        self.assertEqual(tested_line, expected_line)
+
+    def test_get_metadata_line_with_eggnog(self):
+        gene = GeneWithEggNOGFactory()
+        expected_items = [
+            gene.gene_id, gene.name, gene.source, gene.length,
+            '', '', '',
+            '', gene.functions.all()[0].function_id,
+        ]
+        expected_line = ','.join([str(item) for item in expected_items])
+        # Make test with method from GeneViewSet
+        viewset = GeneViewSetMock()
+        tested_line = viewset._get_metadata_line(gene)
+        self.assertEqual(tested_line, expected_line)

backend/metagenedb/apps/catalog/factory/__init__.py

 from .function import EggNOGFactory, FunctionFactory, KeggOrthologyFactory  # noqa
-from .gene import GeneFactory, GeneWithEggNOGFactory, GeneWithKeggFactory  # noqa
+from .gene import GeneFactory, GeneWithEggNOGFactory, GeneWithKeggFactory, GeneWithTaxonomyFactory  # noqa
 from .taxonomy import TaxonomyFactory  # noqa

backend/metagenedb/apps/catalog/factory/gene.py

@@ -11,7 +11,7 @@ from .taxonomy import TaxonomyFactory
 
 faker = Factory.create()
 
-SELECTED_SOURCE = [i[0] for i in models.Function.SOURCE_CHOICES]
+GENE_SOURCES = [i[0] for i in models.Gene.SOURCE_CHOICES]
 
 
 class GeneFactory(DjangoModelFactory):
@@ -21,6 +21,10 @@ class GeneFactory(DjangoModelFactory):
     gene_id = FuzzyLowerText(prefix='gene-', length=15)
     name = fuzzy.FuzzyText(prefix='name-', length=15)
     length = fuzzy.FuzzyInteger(200, 10000)
+    source = fuzzy.FuzzyChoice(GENE_SOURCES)
+
+
+class GeneWithTaxonomyFactory(GeneFactory):
     taxonomy = SubFactory(TaxonomyFactory)
 
 

backend/metagenedb/apps/catalog/operations/test_statistics.py

@@ -2,7 +2,7 @@ from rest_framework.test import APITestCase
 
 from metagenedb.common.utils.color_generator import generate_color_code
 from metagenedb.apps.catalog.factory import (
-    GeneFactory, GeneWithEggNOGFactory, GeneWithKeggFactory, TaxonomyFactory
+    GeneFactory, GeneWithEggNOGFactory, GeneWithKeggFactory, GeneWithTaxonomyFactory, TaxonomyFactory
 )
 
 from .statistics import GeneStatistics, GeneLengthDistribution
@@ -116,7 +116,7 @@ class TestCounts(BaseTestGeneStatistics):
 
         They all have a taxonomy
         """
-        cls.genes_no_function = GeneFactory.create_batch(5)
+        cls.genes_no_function = GeneWithTaxonomyFactory.create_batch(5)
         cls.keggs = GeneWithKeggFactory.create_batch(5)
         cls.eggnogs = GeneWithEggNOGFactory.create_batch(10)
 
@@ -124,7 +124,7 @@ class TestCounts(BaseTestGeneStatistics):
         self.assertEqual(self.gene_stats.count_all(), 20)
 
     def test_has_taxonomy(self):
-        self.assertEqual(self.gene_stats.count_has_taxonomy(), 20)
+        self.assertEqual(self.gene_stats.count_has_taxonomy(), 5)
 
     def test_count_has_function(self):
         self.assertEqual(self.gene_stats.count_has_function(), 15)
@@ -139,7 +139,7 @@ class TestCounts(BaseTestGeneStatistics):
         self.assertEqual(self.gene_stats.count_has_function(source='wrong_source'), 0)
 
     def test_count_has_taxonomy_has_function(self):
-        self.assertEqual(self.gene_stats.count_has_function_has_taxonomy(), 15)
+        self.assertEqual(self.gene_stats.count_has_function_has_taxonomy(), 0)
 
 
 class TestCountWindows(APITestCase):