import json
|
|
import boto3
|
import pytest
|
from io_storages.models import S3ImportStorage
|
from moto import mock_s3
|
from projects.tests.factories import ProjectFactory
|
from rest_framework.test import APIClient
|
from tests.conftest import set_feature_flag_envvar # noqa: F401
|
|
|
@pytest.mark.django_db
|
class TestStoragePredictionValidation:
|
"""Test prediction validation in cloud storage imports."""
|
|
@pytest.fixture
|
def project(self):
|
"""Create a project with a label config for prediction validation."""
|
return ProjectFactory(
|
label_config="""
|
<View>
|
<Text name="text" value="$text"/>
|
<Choices name="sentiment" toName="text">
|
<Choice value="positive"/>
|
<Choice value="negative"/>
|
</Choices>
|
</View>
|
"""
|
)
|
|
@pytest.fixture
|
def api_client(self):
|
"""Create API client for testing."""
|
return APIClient()
|
|
def test_storage_import_with_valid_prediction(self, project, api_client, set_feature_flag_envvar):
|
"""Test that storage import accepts valid predictions."""
|
# Setup API client
|
api_client.force_authenticate(user=project.created_by)
|
|
# Create valid task data with prediction
|
valid_task_data = {
|
'data': {'text': 'This is a positive review'},
|
'predictions': [
|
{
|
'result': [
|
{
|
'from_name': 'sentiment',
|
'to_name': 'text',
|
'type': 'choices',
|
'value': {'choices': ['positive']},
|
}
|
],
|
'score': 0.95,
|
'model_version': 'v1.0',
|
}
|
],
|
}
|
|
with mock_s3():
|
# Setup S3 bucket and test data
|
s3 = boto3.client('s3', region_name='us-east-1')
|
bucket_name = 'pytest-s3-prediction-validation'
|
s3.create_bucket(Bucket=bucket_name)
|
|
# Put valid test data into S3
|
s3.put_object(Bucket=bucket_name, Key='valid_prediction.json', Body=json.dumps([valid_task_data]))
|
|
# Create storage and sync
|
storage = S3ImportStorage(
|
project=project,
|
bucket=bucket_name,
|
aws_access_key_id='example',
|
aws_secret_access_key='example',
|
use_blob_urls=False,
|
)
|
storage.save()
|
storage.sync()
|
|
# Verify task was created
|
tasks_response = api_client.get(f'/api/tasks?project={project.id}')
|
assert tasks_response.status_code == 200
|
tasks = tasks_response.json()['tasks']
|
assert len(tasks) == 1
|
|
# Verify prediction was created
|
predictions_response = api_client.get(f'/api/predictions?task={tasks[0]["id"]}')
|
assert predictions_response.status_code == 200
|
predictions = predictions_response.json()
|
assert len(predictions) == 1
|
|
def test_storage_import_with_invalid_prediction(self, project, api_client, set_feature_flag_envvar):
|
"""Test that storage import rejects invalid predictions."""
|
# Setup API client
|
api_client.force_authenticate(user=project.created_by)
|
|
# Create invalid task data with prediction (wrong from_name)
|
invalid_task_data = {
|
'data': {'text': 'This is a positive review'},
|
'predictions': [
|
{
|
'result': [
|
{
|
'from_name': 'nonexistent_tag', # Invalid from_name
|
'to_name': 'text',
|
'type': 'choices',
|
'value': {'choices': ['positive']},
|
}
|
],
|
'score': 0.95,
|
'model_version': 'v1.0',
|
}
|
],
|
}
|
|
with mock_s3():
|
# Setup S3 bucket and test data
|
s3 = boto3.client('s3', region_name='us-east-1')
|
bucket_name = 'pytest-s3-prediction-validation'
|
s3.create_bucket(Bucket=bucket_name)
|
|
# Put invalid test data into S3
|
s3.put_object(Bucket=bucket_name, Key='invalid_prediction.json', Body=json.dumps([invalid_task_data]))
|
|
# Create storage and sync
|
storage = S3ImportStorage(
|
project=project,
|
bucket=bucket_name,
|
aws_access_key_id='example',
|
aws_secret_access_key='example',
|
use_blob_urls=False,
|
)
|
storage.save()
|
storage.sync()
|
|
# Verify task was NOT created due to validation failure
|
tasks_response = api_client.get(f'/api/tasks?project={project.id}')
|
assert tasks_response.status_code == 200
|
tasks = tasks_response.json()['tasks']
|
assert len(tasks) == 0 # No tasks should be created when predictions are invalid
|
