|
1 | 1 | # -*- coding: utf-8 -*-
|
2 | 2 |
|
| 3 | +import re |
3 | 4 | import sys
|
4 | 5 | import unittest
|
5 | 6 | import StringIO
|
6 | 7 |
|
7 | 8 | from flask import Flask
|
8 |
| -from flask.ext.script import Command, Manager, InvalidCommand, Option |
| 9 | +from flask.ext.script import Command, Manager, Option, prompt, prompt_bool |
| 10 | + |
| 11 | + |
| 12 | +class Catcher(object): |
| 13 | + """Helper decorator to test raw_input.""" |
| 14 | + ## see: http://stackoverflow.com/questions/13480632/python-stringio-selectively-place-data-into-stdin |
| 15 | + |
| 16 | + def __init__(self, handler): |
| 17 | + self.handler = handler |
| 18 | + self.inputs = [] |
| 19 | + |
| 20 | + def __enter__(self): |
| 21 | + self.__stdin = sys.stdin |
| 22 | + self.__stdout = sys.stdout |
| 23 | + sys.stdin = self |
| 24 | + sys.stdout = self |
| 25 | + |
| 26 | + def __exit__(self, type, value, traceback): |
| 27 | + sys.stdin = self.__stdin |
| 28 | + sys.stdout = self.__stdout |
| 29 | + |
| 30 | + def write(self, value): |
| 31 | + self.__stdout.write(value) |
| 32 | + result = self.handler(value) |
| 33 | + if result is not None: |
| 34 | + self.inputs.append(result) |
| 35 | + |
| 36 | + def readline(self): |
| 37 | + return self.inputs.pop() |
| 38 | + |
| 39 | + def getvalue(self): |
| 40 | + return self.__stdout.getvalue() |
| 41 | + |
| 42 | + def truncate(self, pos): |
| 43 | + return self.__stdout.truncate(pos) |
9 | 44 |
|
10 | 45 |
|
11 | 46 | def run(command_line, manager_run, capture_stderr=False):
|
@@ -436,6 +471,87 @@ def test_run_with_default_command(self):
|
436 | 471 | assert code == 0
|
437 | 472 | assert 'OK' in stdout
|
438 | 473 |
|
| 474 | + def test_command_with_prompt(self): |
| 475 | + |
| 476 | + manager = Manager(self.app) |
| 477 | + |
| 478 | + @manager.command |
| 479 | + def hello(): |
| 480 | + print prompt(name='hello') |
| 481 | + |
| 482 | + @Catcher |
| 483 | + def hello_john(msg): |
| 484 | + if re.search("hello", msg): |
| 485 | + return 'john' |
| 486 | + |
| 487 | + with hello_john: |
| 488 | + stdout, code = run('manage.py hello', lambda: manager.run()) |
| 489 | + assert 'hello: john' in stdout |
| 490 | + |
| 491 | + def test_command_with_default_prompt(self): |
| 492 | + |
| 493 | + manager = Manager(self.app) |
| 494 | + |
| 495 | + @manager.command |
| 496 | + def hello(): |
| 497 | + print prompt(name='hello', default='romeo') |
| 498 | + |
| 499 | + @Catcher |
| 500 | + def hello(msg): |
| 501 | + if re.search("hello", msg): |
| 502 | + return '\n' # just hit enter |
| 503 | + |
| 504 | + with hello: |
| 505 | + stdout, code = run('manage.py hello', lambda: manager.run()) |
| 506 | + assert 'hello [romeo]: romeo' in stdout |
| 507 | + |
| 508 | + @Catcher |
| 509 | + def hello_juliette(msg): |
| 510 | + if re.search("hello", msg): |
| 511 | + return 'juliette' |
| 512 | + |
| 513 | + with hello_juliette: |
| 514 | + stdout, code = run('manage.py hello', lambda: manager.run()) |
| 515 | + assert 'hello [romeo]: juliette' in stdout |
| 516 | + |
| 517 | + |
| 518 | + def test_command_with_prompt_bool(self): |
| 519 | + |
| 520 | + manager = Manager(self.app) |
| 521 | + |
| 522 | + @manager.command |
| 523 | + def hello(): |
| 524 | + print prompt_bool(name='correct', default=True, yes_choices=['y'], |
| 525 | + no_choices=['n']) and 'yes' or 'no' |
| 526 | + |
| 527 | + @Catcher |
| 528 | + def correct_default(msg): |
| 529 | + if re.search("correct", msg): |
| 530 | + return '\n' # just hit enter |
| 531 | + |
| 532 | + @Catcher |
| 533 | + def correct_y(msg): |
| 534 | + if re.search("correct", msg): |
| 535 | + return 'y' |
| 536 | + |
| 537 | + @Catcher |
| 538 | + def correct_n(msg): |
| 539 | + if re.search("correct", msg): |
| 540 | + return 'n' |
| 541 | + |
| 542 | + |
| 543 | + with correct_default: |
| 544 | + stdout, code = run('manage.py hello', lambda: manager.run()) |
| 545 | + assert 'correct [y]: yes' in stdout |
| 546 | + |
| 547 | + with correct_y: |
| 548 | + stdout, code = run('manage.py hello', lambda: manager.run()) |
| 549 | + assert 'correct [y]: yes' in stdout |
| 550 | + |
| 551 | + with correct_n: |
| 552 | + stdout, code = run('manage.py hello', lambda: manager.run()) |
| 553 | + assert 'correct [y]: no' in stdout |
| 554 | + |
439 | 555 |
|
440 | 556 | class TestSubManager(unittest.TestCase):
|
441 | 557 |
|
|
0 commit comments