Browse Source

Update CsvWriter, fix #1497

pull/1557/head
Hironsan 3 years ago
parent
commit
bcf6c2a635
4 changed files with 62 additions and 2 deletions
  1. 0
      backend/api/tests/download/__init__.py
  2. 53
      backend/api/tests/download/test_writer.py
  3. 9
      backend/api/views/download/data.py
  4. 2
      backend/api/views/download/writer.py

0
backend/api/tests/download/__init__.py

53
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)

9
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
})

2
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

Loading…
Cancel
Save