122 lines
3.5 KiB
Python
122 lines
3.5 KiB
Python
import pytest
|
|
import swifter
|
|
import pandas as pd
|
|
|
|
from train.mask import FineTrainInstance, strip_parentheses, mask_conditions
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_function():
|
|
return """
|
|
def stream_edit(request, stream_id, response_format="html"):
|
|
"Stream edit page"
|
|
user = request.user.profile
|
|
stream = get_object_or_404(MessageStream, pk=stream_id)
|
|
if not request.user.profile.has_permission(stream, mode="w"):
|
|
return user_denied(
|
|
request,
|
|
message="You don't have access to this Stream",
|
|
response_format=response_format,
|
|
)
|
|
if request.POST:
|
|
if "cancel" not in request.POST:
|
|
form = MessageStreamForm(user, request.POST, instance=stream)
|
|
if form.is_valid():
|
|
stream = form.save()
|
|
return HttpResponseRedirect(
|
|
reverse("messaging_stream_view", args=[stream.id])
|
|
)
|
|
else:
|
|
return HttpResponseRedirect(
|
|
reverse("messaging_stream_view", args=[stream.id])
|
|
)
|
|
else:
|
|
form = MessageStreamForm(user, instance=stream)
|
|
context = _get_default_context(request)
|
|
context.update({"form": form, "stream": stream})
|
|
return render_to_response(
|
|
"messaging/stream_edit",
|
|
context,
|
|
context_instance=RequestContext(request),
|
|
response_format=response_format,
|
|
)
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_function_with_error():
|
|
return """
|
|
def ciao_mamma():
|
|
if 1 > 2:
|
|
print("ciao")
|
|
else if 1 < 2:
|
|
print("ok")
|
|
else:
|
|
return
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_dataframe():
|
|
data = {'source': ['if x > 0: pass', 'if (a and b) or c: pass']}
|
|
df = pd.DataFrame(data)
|
|
|
|
return df
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_dataframe_usi():
|
|
data = {'input_method': ['<fill-in> pass', '<fill-in> pass'], 'target_block': ['if x > 0 :', 'if (a and b) or c :']}
|
|
df = pd.DataFrame(data)
|
|
|
|
return df
|
|
|
|
|
|
def test_mask_does_not_crash(sample_function):
|
|
instances = FineTrainInstance.from_function(sample_function)
|
|
assert len(instances) == 4
|
|
|
|
|
|
def test_mask_with_syntax_error(sample_function_with_error):
|
|
instances = FineTrainInstance.from_function(sample_function_with_error)
|
|
assert instances == []
|
|
|
|
|
|
def test_strip_parentheses_balanced():
|
|
balanced = '("ok")'
|
|
stripped = strip_parentheses(balanced)
|
|
|
|
assert "(" not in stripped and ")" not in stripped
|
|
assert stripped == '"ok"'
|
|
|
|
|
|
def test_strip_parentheses_unbalanced():
|
|
balanced = '("ok"))'
|
|
stripped = strip_parentheses(balanced)
|
|
|
|
assert balanced == stripped
|
|
|
|
|
|
def test_mask_conditions(sample_dataframe):
|
|
result_df = mask_conditions(sample_dataframe, kind='test')
|
|
|
|
assert len(result_df) == 2
|
|
assert 'masked_code' in result_df.columns
|
|
assert 'ground_truth' in result_df.columns
|
|
assert '<mask>' in result_df['masked_code'].iloc[0]
|
|
assert '<mask>' in result_df['masked_code'].iloc[1]
|
|
assert result_df['ground_truth'].iloc[0] == 'x > 0'
|
|
assert result_df['ground_truth'].iloc[1] == 'a and b or c'
|
|
|
|
|
|
def test_mask_conditions_usi(sample_dataframe_usi):
|
|
result_df = mask_conditions(sample_dataframe_usi, kind='test_usi')
|
|
|
|
print(result_df)
|
|
assert len(result_df) == 2
|
|
assert 'masked_code' in result_df.columns
|
|
assert 'ground_truth' in result_df.columns
|
|
assert '<mask>' in result_df['masked_code'].iloc[0]
|
|
assert '<mask>' in result_df['masked_code'].iloc[1]
|
|
assert result_df['ground_truth'].iloc[0] == 'x > 0'
|
|
assert result_df['ground_truth'].iloc[1] == 'a and b or c'
|