From bcf6c2a635e024d22144922c9de550a6a8580f2a Mon Sep 17 00:00:00 2001
From: Hironsan <light.tree.1.13@gmail.com>
Date: Thu, 11 Nov 2021 10:36:19 +0900
Subject: [PATCH] Update CsvWriter, fix #1497

---
 backend/api/tests/download/__init__.py    |  0
 backend/api/tests/download/test_writer.py | 53 +++++++++++++++++++++++
 backend/api/views/download/data.py        |  9 +++-
 backend/api/views/download/writer.py      |  2 +-
 4 files changed, 62 insertions(+), 2 deletions(-)
 create mode 100644 backend/api/tests/download/__init__.py
 create mode 100644 backend/api/tests/download/test_writer.py

diff --git a/backend/api/tests/download/__init__.py b/backend/api/tests/download/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/backend/api/tests/download/test_writer.py b/backend/api/tests/download/test_writer.py
new file mode 100644
index 00000000..2c48cb92
--- /dev/null
+++ b/backend/api/tests/download/test_writer.py
@@ -0,0 +1,53 @@
+import unittest
+from unittest.mock import call, patch
+
+from ...views.download.data import Record
+from ...views.download.writer import CsvWriter
+
+
+class TestCSVWriter(unittest.TestCase):
+
+    def setUp(self):
+        self.records = [
+            Record(id=0, data='exampleA', label=['labelA'], user='admin', metadata={'hidden': 'secretA'}),
+            Record(id=1, data='exampleB', label=['labelB'], user='admin', metadata={'hidden': 'secretB'}),
+            Record(id=2, data='exampleC', label=['labelC'], user='admin', metadata={'meta': 'secretC'})
+        ]
+
+    def test_create_header(self):
+        writer = CsvWriter('.')
+        header = writer.create_header(self.records)
+        expected = ['id', 'data', 'label', 'hidden', 'meta']
+        self.assertEqual(header, expected)
+
+    def test_create_line(self):
+        writer = CsvWriter('.')
+        record = self.records[0]
+        line = writer.create_line(record)
+        expected = {
+            'id': record.id,
+            'data': record.data,
+            'label': record.label[0],
+            'hidden': 'secretA'
+        }
+        self.assertEqual(line, expected)
+
+    @patch('os.remove')
+    @patch('zipfile.ZipFile')
+    @patch('csv.DictWriter.writerow')
+    @patch('builtins.open')
+    def test_dump(self, mock_open_file, csv_io, zip_io, mock_remove_file):
+        writer = CsvWriter('.')
+        writer.write(self.records)
+
+        self.assertEqual(mock_open_file.call_count, 1)
+        mock_open_file.assert_called_with('./admin.csv', mode='a', encoding='utf-8')
+
+        self.assertEqual(csv_io.call_count, len(self.records) + 1)  # +1 is for a header
+        calls = [
+            call({'id': 'id', 'data': 'data', 'label': 'label', 'hidden': 'hidden', 'meta': 'meta'}),
+            call({'id': 0, 'data': 'exampleA', 'label': 'labelA', 'hidden': 'secretA'}),
+            call({'id': 1, 'data': 'exampleB', 'label': 'labelB', 'hidden': 'secretB'}),
+            call({'id': 2, 'data': 'exampleC', 'label': 'labelC', 'meta': 'secretC'})
+        ]
+        csv_io.assert_has_calls(calls)
diff --git a/backend/api/views/download/data.py b/backend/api/views/download/data.py
index 68978184..79bfe7e8 100644
--- a/backend/api/views/download/data.py
+++ b/backend/api/views/download/data.py
@@ -1,3 +1,4 @@
+import json
 from typing import Any, Dict, List
 
 
@@ -16,4 +17,10 @@ class Record:
         self.metadata = metadata
 
     def __str__(self):
-        return f'{self.data}\t{self.label}'
+        return json.dumps({
+            'id': self.id,
+            'data': self.data,
+            'label': self.label,
+            'user': self.user,
+            'metadata': self.metadata
+        })
diff --git a/backend/api/views/download/writer.py b/backend/api/views/download/writer.py
index dd0e88b4..5de1264e 100644
--- a/backend/api/views/download/writer.py
+++ b/backend/api/views/download/writer.py
@@ -90,7 +90,7 @@ class CsvWriter(BaseWriter):
 
     def create_header(self, records: List[Record]) -> Iterable[str]:
         header = ['id', 'data', 'label']
-        header += list(itertools.chain(*[r.metadata.keys() for r in records]))
+        header += sorted(set(itertools.chain(*[r.metadata.keys() for r in records])))
         return header