"""This file and its contents are licensed under the Apache License 2.0. Please see the included NOTICE for copyright information and LICENSE for a copy of the license.
|
"""
|
import copy
|
import io
|
import zipfile
|
|
import pytest
|
import requests_mock
|
import ujson as json
|
from projects.models import Project
|
from rest_framework.authtoken.models import Token
|
from tasks.models import Annotation, Prediction, Task
|
|
|
def post_data_as_format(setup, format_type, body, archive, multiply_files):
|
# post as data
|
if format_type == 'json_data':
|
return setup.post(setup.urls.task_bulk, data=body, content_type='application/json')
|
|
# post as files
|
if format_type == 'json_file':
|
files = {f'upload_file{i}.json': io.StringIO(body) for i in range(0, multiply_files)}
|
elif format_type == 'csv_file':
|
files = {f'upload_file{i}.csv': io.StringIO(body) for i in range(0, multiply_files)}
|
elif format_type == 'tsv_file':
|
files = {f'upload_file{i}.tsv': io.StringIO(body) for i in range(0, multiply_files)}
|
elif format_type == 'txt_file':
|
files = {f'upload_file{i}.txt': io.StringIO(body) for i in range(0, multiply_files)}
|
else:
|
raise Exception('Incorrect task data format to post')
|
|
# zip: take files below and zip them
|
if 'zip' in archive:
|
file = io.BytesIO()
|
ref = zipfile.ZipFile(file, mode='w', compression=zipfile.ZIP_DEFLATED)
|
[ref.writestr(name, body.read()) for name, body in files.items()]
|
|
ref.close()
|
file.seek(0, 0)
|
files = {'upload_file.zip': file}
|
|
# replicate zip file x2
|
if 'zip_x2' == archive:
|
files.update({'upload_file2.zip': copy.deepcopy(file)})
|
|
return setup.post(setup.urls.task_bulk, files)
|
|
|
@pytest.mark.parametrize('multiply_files', [1, 5])
|
@pytest.mark.parametrize('format_type', ['json_file', 'json_data'])
|
@pytest.mark.parametrize(
|
'tasks, status_code, task_count',
|
[
|
([{'data': {'dialog': 'some'}}], 201, 1),
|
([{'data': {'dialog': 'some'}}] * 10, 201, 10),
|
([{'data': {'another_field': 'some', 'dialog': 'some'}}], 201, 1),
|
([{'data': {'dialog': 123}, 'created_at': 123}], 201, 1),
|
([{'data': {'another_field': 'some'}}] * 10, 400, 0),
|
([{'data': {}}], 400, 0),
|
([{'data': None}], 400, 0),
|
(None, 400, 0),
|
([{'data': 'string'}], 400, 0),
|
([{}, {}], 400, 0),
|
([{}], 400, 0),
|
({}, 400, 0),
|
([], 400, 0),
|
([{'dialog': 'some'}] * 10, 201, 10),
|
({'dialog': 'some'}, 201, 1),
|
([{'dialog': 'some', 'second_field': 123}] * 10, 201, 10),
|
([{'none': 'some', 'second_field': 123}] * 10, 400, 0),
|
],
|
)
|
@pytest.mark.django_db
|
def test_json_task_upload(setup_project_dialog, format_type, tasks, status_code, task_count, multiply_files):
|
"""Upload JSON as file and data with one task to project.
|
Decorator pytest.mark.django_db means it will be clean DB setup_project_dialog for this test.
|
"""
|
if format_type == 'json_data' and multiply_files > 1:
|
pytest.skip('Senseless parameter combination')
|
|
r = post_data_as_format(setup_project_dialog, format_type, json.dumps(tasks), 'none', multiply_files)
|
print(f'Create json {format_type} tasks result:', r.content)
|
assert r.status_code == status_code, f'Upload tasks failed. Response data: {r.data}'
|
assert Task.objects.filter(project=setup_project_dialog.project.id).count() == task_count * multiply_files
|
|
|
@pytest.mark.parametrize(
|
'tasks, status_code, task_count, annotation_count',
|
[
|
([{'data': {'dialog': 'Test'}, 'annotations': [{'result': [{'id': '123'}]}]}] * 10, 201, 10, 10),
|
(
|
[{'data': {'dialog': 'Test'}, 'annotations': [{'result': [{'id': '123'}], 'ground_truth': True}]}],
|
201,
|
1,
|
1,
|
),
|
([{'data': {'dialog': 'Test'}, 'annotations': [{'result': '123'}]}], 400, 0, 0),
|
([{'data': {'dialog': 'Test'}, 'meta': 'test'}] * 10, 400, 0, 0),
|
([{'data': {'dialog': 'Test'}, 'annotations': 'test'}] * 10, 400, 0, 0),
|
([{'data': {'dialog': 'Test'}, 'annotations': [{'trash': '123'}]}] * 10, 400, 0, 0),
|
],
|
)
|
@pytest.mark.django_db
|
def test_json_task_annotation_and_meta_upload(setup_project_dialog, tasks, status_code, task_count, annotation_count):
|
"""Upload JSON task with annotation to project"""
|
format_type = 'json_file'
|
multiply_files = 1
|
|
r = post_data_as_format(setup_project_dialog, format_type, json.dumps(tasks), 'none', multiply_files)
|
print('Create json tasks with annotations result:', r.content)
|
assert r.status_code == status_code, 'Upload one task with annotation failed'
|
|
# tasks
|
tasks_db = Task.objects.filter(project=setup_project_dialog.project.id)
|
assert tasks_db.count() == task_count * multiply_files
|
for task in tasks_db:
|
assert task.is_labeled, 'Task should be labeled'
|
|
# annotations
|
annotations = Annotation.objects.filter(task__project=setup_project_dialog.project.id)
|
assert annotations.count() == annotation_count * multiply_files
|
for i, annotation in enumerate(annotations):
|
assert annotation.ground_truth
|
|
|
@pytest.mark.parametrize(
|
'tasks, status_code, task_count, prediction_count',
|
[
|
(
|
[
|
{
|
'data': {'dialog': 'Test'},
|
'predictions': [
|
{
|
'result': [
|
{
|
'id': '123',
|
'from_name': 'answer',
|
'to_name': 'dialog',
|
'type': 'textarea',
|
'value': {'text': ['Test prediction']},
|
}
|
],
|
'model_version': 'test',
|
}
|
],
|
}
|
],
|
201,
|
1,
|
1,
|
),
|
([{'data': {'dialog': 'Test'}, 'predictions': [{'WRONG_FIELD': '123'}]}], 400, 0, 0),
|
],
|
)
|
@pytest.mark.django_db
|
def test_json_task_predictions(setup_project_dialog, tasks, status_code, task_count, prediction_count):
|
"""Upload JSON task with predictions to project"""
|
r = post_data_as_format(setup_project_dialog, 'json_file', json.dumps(tasks), 'none', 1)
|
assert r.status_code == status_code, 'Upload one task with prediction failed'
|
|
# predictions
|
predictions = Prediction.objects.filter(project=setup_project_dialog.project.id)
|
assert predictions.count() == prediction_count
|
for i, predictions in enumerate(predictions):
|
assert predictions.model_version == 'test'
|
|
|
@pytest.mark.parametrize('multiply_files', [1, 5])
|
@pytest.mark.parametrize('archive', ['none'])
|
@pytest.mark.parametrize('format_type', ['json_file'])
|
@pytest.mark.parametrize(
|
'tasks, status_code, task_count, annotation_count',
|
[
|
(
|
[{'data': {'dialog': 'Test'}, 'annotations': [{'result': [{'id': '123'}]}, {'result': [{'id': '456'}]}]}]
|
* 10,
|
201,
|
10,
|
20,
|
),
|
([{'data': {'dialog': 'Test'}, 'annotations': [{'trash': '123'}]}] * 10, 400, 0, 0),
|
],
|
)
|
@pytest.mark.django_db
|
def test_archives(
|
setup_project_dialog, format_type, tasks, status_code, task_count, annotation_count, archive, multiply_files
|
):
|
"""Upload JSON task with annotation to project"""
|
multiplier = (2 if 'zip_x2' == archive else 1) * multiply_files
|
|
r = post_data_as_format(setup_project_dialog, format_type, json.dumps(tasks), archive, multiply_files)
|
print('Create json tasks with annotations result:', r.content)
|
assert r.status_code == status_code, 'Upload one task with annotation failed'
|
|
# tasks
|
tasks = Task.objects.filter(project=setup_project_dialog.project.id)
|
assert tasks.count() == task_count * multiplier
|
for task in tasks:
|
assert task.is_labeled, 'Task should be labeled'
|
|
# annotations
|
annotations = Annotation.objects.filter(task__project=setup_project_dialog.project.id)
|
assert annotations.count() == annotation_count * multiplier
|
for annotation in annotations:
|
assert annotation.ground_truth
|
|
|
@pytest.mark.parametrize('multiply_files', [1, 5])
|
@pytest.mark.parametrize('archive', ['none'])
|
@pytest.mark.parametrize('format_type', ['csv_file', 'tsv_file'])
|
@pytest.mark.parametrize(
|
'tasks, status_code, task_count',
|
[
|
('dialog,second\ndialog 1,second 1\ndialog 2,second 2', 201, 2),
|
('dialog,second,class\ndialog 1, second 2, class 1', 201, 1),
|
('here_is_error_in_column_count,second\ndialog 1, second 1, class 1', 400, 0),
|
('empty_rows\n', 400, 0),
|
('', 400, 0),
|
],
|
)
|
@pytest.mark.django_db
|
def test_csv_tsv_task_upload(
|
setup_project_dialog, format_type, tasks, status_code, task_count, archive, multiply_files
|
):
|
"""Upload CSV/TSV with one task to project"""
|
multiplier = (2 if 'zip_x2' == archive else 1) * multiply_files
|
|
tasks = tasks if format_type == 'csv_file' else tasks.replace(',', '\t') # prepare tsv file from csv
|
r = post_data_as_format(setup_project_dialog, format_type, tasks, archive, multiply_files)
|
print(f'Create {format_type} tasks result:', r.content)
|
|
assert r.status_code == status_code, f'Upload one task {format_type} failed. Response data: {r.data}'
|
assert Task.objects.filter(project=setup_project_dialog.project.id).count() == task_count * multiplier
|
|
|
@pytest.mark.parametrize('multiply_files', [1, 5])
|
@pytest.mark.parametrize('format_type', ['txt_file'])
|
@pytest.mark.parametrize('tasks, status_code, task_count', [('my text 1\nmy text 2\nmy text 3', 201, 3), ('', 400, 0)])
|
@pytest.mark.django_db
|
def test_txt_task_upload(setup_project_dialog, format_type, tasks, status_code, task_count, multiply_files):
|
"""Upload CSV/TSV with one task to project"""
|
multiplier = multiply_files
|
|
r = post_data_as_format(setup_project_dialog, format_type, tasks, 'none', multiply_files)
|
print(f'Create {format_type} tasks result:', r.content)
|
|
assert r.status_code == status_code, f'Upload one task {format_type} failed. Response data: {r.data}'
|
assert Task.objects.filter(project=setup_project_dialog.project.id).count() == task_count * multiplier
|
|
|
@pytest.mark.parametrize(
|
'tasks, status_code, task_count, max_duration',
|
[([{'data': {'dialog': 'Test'}, 'annotations': [{'result': [{'id': '123'}]}]}] * 1000, 201, 1000, 30)],
|
)
|
@pytest.mark.django_db
|
def test_upload_duration(setup_project_dialog, tasks, status_code, task_count, max_duration):
|
"""Upload JSON task with annotation to project"""
|
r = post_data_as_format(setup_project_dialog, 'json_data', json.dumps(tasks), 'none', 1)
|
print('Create json tasks with annotations result:', r.content)
|
assert r.status_code == status_code, ('Upload one task with annotation failed', r.content)
|
|
# tasks
|
tasks = Task.objects.filter(project=setup_project_dialog.project.id)
|
assert tasks.count() == task_count
|
for task in tasks:
|
assert task.is_labeled, 'Task should be labeled'
|
|
# check max duration
|
result = json.loads(r.content)
|
assert result['duration'] < max_duration, 'Max duration of adding tasks is exceeded'
|
|
|
@pytest.mark.parametrize(
|
'tasks, status_code, task_count',
|
[([{'data': {'dialog': 'Test'}, 'annotations': [{'result': [{'id': '123'}]}]}] * 100, 201, 100)],
|
)
|
@pytest.mark.django_db
|
def test_url_upload(mocker, setup_project_dialog, tasks, status_code, task_count):
|
"""Upload tasks from URL"""
|
with requests_mock.Mocker(real_http=True) as m:
|
url = 'http://localhost:8111/test.json'
|
m.get(url, text=json.dumps(tasks), headers={'Content-Length': '100'})
|
r = setup_project_dialog.post(
|
setup_project_dialog.urls.task_bulk, data='url=' + url, content_type='application/x-www-form-urlencoded'
|
)
|
assert r.status_code == status_code, 'Upload URL failed: ' + str(r.content)
|
|
# tasks
|
tasks = Task.objects.filter(project=setup_project_dialog.project.id)
|
assert tasks.count() == task_count
|
for task in tasks:
|
assert task.is_labeled, 'Task should be labeled since annotation is ground_truth'
|
|
|
@pytest.mark.parametrize(
|
'tasks, status_code, task_count, bad_token',
|
[([{'dialog': 'Test'}] * 1, 201, 1, False), ([{'dialog': 'Test'}] * 1, 401, 0, True)],
|
)
|
@pytest.mark.django_db
|
def test_upload_with_token(setup_project_for_token, tasks, status_code, task_count, bad_token):
|
"""Upload with Django Token"""
|
setup = setup_project_for_token
|
token = Token.objects.get(user=setup.user)
|
token = 'Token ' + str(token)
|
broken_token = 'Token broken'
|
data = setup.project_config
|
data['organization_pk'] = setup.org.pk
|
r = setup.post(setup.urls.project_create, data=data, HTTP_AUTHORIZATION=token)
|
print('Project create with status code:', r.status_code, r.content)
|
assert r.status_code == 201, 'Create project result should be redirect to the next page: ' + str(r.content)
|
|
project = Project.objects.filter(title=setup.project_config['title']).first()
|
setup.urls.set_project(project.pk)
|
|
r = setup.post(
|
setup.urls.task_bulk,
|
data=json.dumps(tasks),
|
content_type='application/json',
|
HTTP_AUTHORIZATION=broken_token if bad_token else token,
|
)
|
assert r.status_code == status_code, 'Create json tasks result: ' + str(r.content)
|
|
# tasks
|
tasks = Task.objects.filter(project=project.id)
|
assert tasks.count() == task_count
|
