doccano/backend/data_export/tests/test_writer.py

import json
import unittest
from unittest.mock import call, patch

from ..pipeline.data import Record
from ..pipeline.writers import CsvWriter, IntentAndSlotWriter


class TestCSVWriter(unittest.TestCase):
    def setUp(self):
        self.records = [
            Record(data_id=0, data="exampleA", label=["labelA"], user="admin", metadata={"hidden": "secretA"}),
            Record(data_id=1, data="exampleB", label=["labelB"], user="admin", metadata={"hidden": "secretB"}),
            Record(data_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)

    def test_label_order(self):
        writer = CsvWriter(".")
        record1 = Record(data_id=0, data="", label=["labelA", "labelB"], user="", metadata={})
        record2 = Record(data_id=0, data="", label=["labelB", "labelA"], user="", metadata={})
        line1 = writer.create_line(record1)
        line2 = writer.create_line(record2)
        expected = "labelA#labelB"
        self.assertEqual(line1["label"], expected)
        self.assertEqual(line2["label"], 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)


class TestIntentWriter(unittest.TestCase):
    def setUp(self):
        self.record = Record(
            data_id=0,
            data="exampleA",
            label={"cats": ["positive"], "entities": [(0, 1, "LOC")]},
            user="admin",
            metadata={},
        )

    def test_create_line(self):
        writer = IntentAndSlotWriter(".")
        actual = writer.create_line(self.record)
        expected = {
            "id": self.record.id,
            "text": self.record.data,
            "cats": ["positive"],
            "entities": [[0, 1, "LOC"]],
        }
        self.assertEqual(json.loads(actual), expected)