doccano/backend/auto_labeling/tests/test_views.py

import pathlib
from unittest.mock import patch

from auto_labeling_pipeline.mappings import AmazonComprehendSentimentTemplate
from auto_labeling_pipeline.models import RequestModelFactory
from model_mommy import mommy
from rest_framework import status
from rest_framework.reverse import reverse

from api.models import DOCUMENT_CLASSIFICATION, IMAGE_CLASSIFICATION
from api.tests.api.utils import (CRUDMixin, make_auto_labeling_config, make_doc, make_image,
                                 prepare_project)

data_dir = pathlib.Path(__file__).parent / 'data'


class TestConfigParameter(CRUDMixin):

    def setUp(self):
        self.project = prepare_project(task=DOCUMENT_CLASSIFICATION)
        self.data = {
            'model_name': 'GCP Entity Analysis',
            'model_attrs': {'key': 'hoge', 'type': 'PLAIN_TEXT', 'language': 'en'},
            'text': 'example'
        }
        self.url = reverse(viewname='auto_labeling_parameter_testing', args=[self.project.item.id])

    @patch('auto_labeling.views.AutoLabelingConfigParameterTest.send_request', return_value={})
    def test_called_with_proper_model(self, mock):
        self.assert_create(self.project.users[0], status.HTTP_200_OK)
        _, kwargs = mock.call_args
        expected = RequestModelFactory.create(self.data['model_name'], self.data['model_attrs'])
        self.assertEqual(kwargs['model'], expected)

    @patch('auto_labeling.views.AutoLabelingConfigParameterTest.send_request', return_value={})
    def test_called_with_text(self, mock):
        self.assert_create(self.project.users[0], status.HTTP_200_OK)
        _, kwargs = mock.call_args
        self.assertEqual(kwargs['example'], self.data['text'])

    @patch('auto_labeling.views.AutoLabelingConfigParameterTest.send_request', return_value={})
    def test_called_with_image(self, mock):
        self.data['text'] = str(data_dir / 'images/1500x500.jpeg')
        self.assert_create(self.project.users[0], status.HTTP_200_OK)
        _, kwargs = mock.call_args
        self.assertEqual(kwargs['example'], self.data['text'])


class TestTemplateMapping(CRUDMixin):

    def setUp(self):
        self.project = prepare_project(task=DOCUMENT_CLASSIFICATION)
        self.data = {
            'response': {
                'Sentiment': 'NEUTRAL',
                'SentimentScore': {
                    'Positive': 0.004438233096152544,
                    'Negative': 0.0005306027014739811,
                    'Neutral': 0.9950305223464966,
                    'Mixed': 5.80838445785048e-7
                }
            },
            'template': AmazonComprehendSentimentTemplate().load()
        }
        self.url = reverse(viewname='auto_labeling_template_test', args=[self.project.item.id])

    def test_template_mapping(self):
        response = self.assert_create(self.project.users[0], status.HTTP_200_OK)
        expected = [{'label': 'NEUTRAL'}]
        self.assertEqual(response.json(), expected)

    def test_json_decode_error(self):
        self.data['template'] = ''
        self.assert_create(self.project.users[0], status.HTTP_400_BAD_REQUEST)


class TestLabelMapping(CRUDMixin):

    def setUp(self):
        self.project = prepare_project(task=DOCUMENT_CLASSIFICATION)
        self.data = {
            'response': [{'label': 'NEGATIVE'}],
            'label_mapping': {'NEGATIVE': 'Negative'}
        }
        self.url = reverse(viewname='auto_labeling_mapping_test', args=[self.project.item.id])

    def test_label_mapping(self):
        response = self.assert_create(self.project.users[0], status.HTTP_200_OK)
        expected = [{'label': 'Negative'}]
        self.assertEqual(response.json(), expected)


class TestConfigCreation(CRUDMixin):

    def setUp(self):
        self.project = prepare_project(task=DOCUMENT_CLASSIFICATION)
        self.data = {
            'model_name': 'Amazon Comprehend Sentiment Analysis',
            'model_attrs': {
                'aws_access_key': 'str',
                'aws_secret_access_key': 'str',
                'region_name': 'us-east-1',
                'language_code': 'en'
            },
            'template': AmazonComprehendSentimentTemplate().load(),
            'label_mapping': {'NEGATIVE': 'Negative'}
        }
        self.url = reverse(viewname='auto_labeling_configs', args=[self.project.item.id])

    def test_create_config(self):
        response = self.assert_create(self.project.users[0], status.HTTP_201_CREATED)
        self.assertEqual(response.data['model_name'], self.data['model_name'])

    def test_list_config(self):
        mommy.make('AutoLabelingConfig', project=self.project.item)
        response = self.assert_fetch(self.project.users[0], status.HTTP_200_OK)
        self.assertEqual(len(response.data), 1)


class TestAutoLabelingText(CRUDMixin):

    def setUp(self):
        self.project = prepare_project(task=DOCUMENT_CLASSIFICATION)
        make_auto_labeling_config(self.project.item)
        self.example = make_doc(self.project.item)
        self.url = reverse(viewname='auto_labeling_annotation', args=[self.project.item.id, self.example.id])

    @patch('auto_labeling.views.execute_pipeline', return_value=[])
    def test_text_task(self, mock):
        self.assert_create(self.project.users[0], status.HTTP_201_CREATED)
        _, kwargs = mock.call_args
        self.assertEqual(kwargs['text'], self.example.text)


class TestAutoLabelingImage(CRUDMixin):

    def setUp(self):
        self.project = prepare_project(task=IMAGE_CLASSIFICATION)
        make_auto_labeling_config(self.project.item)
        filepath = data_dir / 'images/1500x500.jpeg'
        self.example = make_image(self.project.item, str(filepath))
        self.url = reverse(viewname='auto_labeling_annotation', args=[self.project.item.id, self.example.id])

    @patch('auto_labeling.views.execute_pipeline', return_value=[])
    def test_text_task(self, mock):
        self.assert_create(self.project.users[0], status.HTTP_201_CREATED)
        _, kwargs = mock.call_args
        expected = str(self.example.filename)
        self.assertEqual(kwargs['text'], expected)