On reflection, using subprocess for a case like this is overkill and just too inefficient to tolerate. Exec is a much better approach. It allows you to run the student job up to about 100,000 times to give really good stats, rather than the 100 or so that's possible with subprocess. The example below runs it 1000 times which should be plenty for a simple case like this.
If the student code fails within the exec - even if it gets stuck in a loop forcing a timeout - you still get to see the "time limit exceeded" error message.
import sys, io
TOLERATED_DIFFERENCE = 300 # What's a good number?!
prog = """{{ STUDENT_ANSWER | e('py') }}"""
num_great = 0
num_rubbish = 0
message = ''
saved_stdout = sys.stdout
try:
for i_test in range(1000):
sys.stdout = io.StringIO()
exec(prog) # Run the student code
result = sys.stdout.getvalue().rstrip()
if result == 'Great!':
num_great += 1
elif result == 'Rubbish!':
num_rubbish += 1
else:
message = "Sorry, but the output '{}' isn't valid".format(result)
break
if not message:
if abs(num_great - num_rubbish) < TOLERATED_DIFFERENCE:
message = "It worked!"
else:
message = "Output was wrong.\n"
message += "You printed 'Great!' {} times and 'Rubbish! {} times".format(num_great, num_rubbish)
except Exception as err:
message = "I couldn't run your code!\n"
message += str(err)
sys.stdout = saved_stdout
print(message)