Skip to content

Commit

Permalink
tests: Update some integration tests
Browse files Browse the repository at this point in the history
  • Loading branch information
frascuchon committed Mar 5, 2024
1 parent cf7f3ab commit 98e9cca
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 46 deletions.
64 changes: 30 additions & 34 deletions tests/integration/client/feedback/dataset/local/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
import argilla.client.singleton
import datasets
import pytest
from argilla import User, Workspace
from argilla import ResponseSchema, User, Workspace
from argilla.client.feedback.config import DatasetConfig
from argilla.client.feedback.dataset import FeedbackDataset
from argilla.client.feedback.schemas.fields import TextField
Expand All @@ -28,6 +28,7 @@
TermsMetadataProperty,
)
from argilla.client.feedback.schemas.questions import SpanLabelOption, SpanQuestion, TextQuestion
from argilla.client.feedback.schemas.suggestions import SuggestionSchema
from argilla.client.feedback.schemas.records import FeedbackRecord
from argilla.client.feedback.schemas.remote.records import RemoteSuggestionSchema
from argilla.client.feedback.schemas.vector_settings import VectorSettings
Expand All @@ -51,7 +52,9 @@ def test_create_dataset_with_suggestions(argilla_user: "ServerUser") -> None:
records=[
FeedbackRecord(
fields={"text": "this is a text"},
suggestions=[{"question_name": "text", "value": "This is a suggestion"}],
suggestions=[
SuggestionSchema.with_question_value(ds.question_by_name("text"), value="This is a suggestion")
],
)
]
)
Expand All @@ -69,7 +72,7 @@ def test_create_dataset_with_span_questions(argilla_user: "ServerUser") -> None:

ds = FeedbackDataset(
fields=[TextField(name="text")],
questions=[SpanQuestion(name="spans", labels=["label1", "label2"])],
questions=[SpanQuestion(name="spans", field="text", labels=["label1", "label2"])],
)

rg_dataset = ds.push_to_argilla(name="new_dataset")
Expand All @@ -93,7 +96,9 @@ async def test_update_dataset_records_with_suggestions(argilla_user: "ServerUser
assert remote_dataset.records[0].id is not None
assert remote_dataset.records[0].suggestions == ()

remote_dataset.records[0].update(suggestions=[{"question_name": "text", "value": "This is a suggestion"}])
remote_dataset.records[0].update(
suggestions=[SuggestionSchema.with_question_value(ds.question_by_name("text"), value="This is a suggestion")]
)

# TODO: Review this requirement for tests and explain, try to avoid use or at least, document.
await db.refresh(argilla_user, attribute_names=["datasets"])
Expand Down Expand Up @@ -139,6 +144,11 @@ def test_add_records(
assert not dataset.records[0].responses
assert not dataset.records[0].suggestions

question_1 = dataset.question_by_name("question-1")
question_2 = dataset.question_by_name("question-2")
question_3 = dataset.question_by_name("question-3")
question_4 = dataset.question_by_name("question-4")
question_5 = dataset.question_by_name("question-5")
dataset.add_records(
[
FeedbackRecord(
Expand All @@ -148,38 +158,24 @@ def test_add_records(
},
metadata={"unit": "test"},
responses=[
{
"values": {
"question-1": {"value": "answer"},
"question-2": {"value": 0},
"question-3": {"value": "a"},
"question-4": {"value": ["a", "b"]},
"question-5": {"value": [{"rank": 1, "value": "a"}, {"rank": 2, "value": "b"}]},
},
"status": "submitted",
},
ResponseSchema(status="submitted")
.with_question_value(question_1, value="answer")
.with_question_value(question_2, value=0)
.with_question_value(question_3, value="a")
.with_question_value(question_4, value=["a", "b"])
.with_question_value(
question_5,
value=[{"rank": 1, "value": "a"}, {"rank": 2, "value": "b"}],
)
],
suggestions=[
{
"question_name": "question-1",
"value": "answer",
},
{
"question_name": "question-2",
"value": 0,
},
{
"question_name": "question-3",
"value": "a",
},
{
"question_name": "question-4",
"value": ["a", "b"],
},
{
"question_name": "question-5",
"value": [{"rank": 1, "value": "a"}, {"rank": 2, "value": "b"}],
},
SuggestionSchema.with_question_value(question_1, value="answer"),
SuggestionSchema.with_question_value(question_2, value=0),
SuggestionSchema.with_question_value(question_3, value="a"),
SuggestionSchema.with_question_value(question_4, value=["a", "b"]),
SuggestionSchema.with_question_value(
question_5, value=[{"rank": 1, "value": "a"}, {"rank": 2, "value": "b"}]
),
],
external_id="test-id",
),
Expand Down
23 changes: 11 additions & 12 deletions tests/integration/client/feedback/dataset/remote/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,21 +11,18 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest
import random
from datetime import datetime
from typing import TYPE_CHECKING, Any, List, Tuple, Type
from sqlalchemy.ext.asyncio import AsyncSession
from typing import Any, List, TYPE_CHECKING, Tuple, Type
from uuid import UUID

import argilla as rg
import argilla.client.singleton
import numpy
import pytest
from argilla import (
FeedbackRecord,
)
from argilla import FeedbackRecord, SuggestionSchema
from argilla.client.feedback.dataset import FeedbackDataset
from argilla.client.feedback.dataset.remote.dataset import RemoteFeedbackDataset
from argilla.client.feedback.schemas import SuggestionSchema
from argilla.client.feedback.schemas.fields import TextField
from argilla.client.feedback.schemas.metadata import (
FloatMetadataProperty,
Expand All @@ -51,8 +48,6 @@
from argilla.client.workspaces import Workspace
from argilla_server.models import User as ServerUser
from argilla_server.settings import settings
from sqlalchemy.ext.asyncio import AsyncSession

from tests.factories import (
DatasetFactory,
RecordFactory,
Expand Down Expand Up @@ -166,11 +161,12 @@ async def test_update_records_with_suggestions(
)

remote = test_dataset_with_metadata_properties.push_to_argilla(name="test_dataset", workspace=ws)
question = remote.question_by_name("question")

records = []
for record in remote:
record.suggestions = [
SuggestionSchema(question_name="question", value=f"Hello world! for {record.fields['text']}")
SuggestionSchema.with_question_value(question, value=f"Hello world! for {record.fields['text']}")
]
records.append(record)

Expand All @@ -192,14 +188,17 @@ async def test_update_records_with_empty_list_of_suggestions(
ws = rg.Workspace.create(name="test-workspace")

remote = test_dataset_with_metadata_properties.push_to_argilla(name="test_dataset", workspace=ws)
question = remote.question_by_name("question")

remote.add_records(
[
FeedbackRecord(
fields={"text": "Hello world!"}, suggestions=[{"question_name": "question", "value": "test"}]
fields={"text": "Hello world!"},
suggestions=[SuggestionSchema.with_question_value(question, value="test")],
),
FeedbackRecord(
fields={"text": "Another record"}, suggestions=[{"question_name": "question", "value": "test"}]
fields={"text": "Another record"},
suggestions=[SuggestionSchema.with_question_value(question, value="test")],
),
]
)
Expand Down

0 comments on commit 98e9cca

Please sign in to comment.