<s> import datetime <EOL> import os <EOL> try : <EOL> from StringIO import StringIO <EOL> except ImportError : <EOL> from io import StringIO <EOL> import sys <EOL> from peewee import * <EOL> from pwiz import * <EOL> from playhouse . tests . base import database_initializer <EOL> from playhouse . tests . base import mock <EOL> from playhouse . tests . base import PeeweeTestCase <EOL> from playhouse . tests . base import skip_if <EOL> db = database_initializer . get_database ( '<STR_LIT>' ) <EOL> class BaseModel ( Model ) : <EOL> class Meta : <EOL> database = db <EOL> class User ( BaseModel ) : <EOL> username = CharField ( primary_key = True ) <EOL> id = IntegerField ( default = <NUM_LIT:0> ) <EOL> class Note ( BaseModel ) : <EOL> user = ForeignKeyField ( User ) <EOL> text = TextField ( index = True ) <EOL> data = IntegerField ( default = <NUM_LIT:0> ) <EOL> misc = IntegerField ( default = <NUM_LIT:0> ) <EOL> class Meta : <EOL> indexes = ( <EOL> ( ( '<STR_LIT:user>' , '<STR_LIT:text>' ) , True ) , <EOL> ( ( '<STR_LIT:user>' , '<STR_LIT:data>' , '<STR_LIT>' ) , False ) , <EOL> ) <EOL> class Category ( BaseModel ) : <EOL> name = CharField ( unique = True ) <EOL> parent = ForeignKeyField ( '<STR_LIT>' , null = True ) <EOL> class capture_output ( object ) : <EOL> def __enter__ ( self ) : <EOL> self . _stdout = sys . stdout <EOL> sys . stdout = self . _buffer = StringIO ( ) <EOL> return self <EOL> def __exit__ ( self , * args ) : <EOL> self . data = self . _buffer . getvalue ( ) <EOL> sys . stdout = self . _stdout <EOL> EXPECTED = """<STR_LIT>""" . strip ( ) <EOL> EXPECTED_ORDERED = """<STR_LIT>""" . strip ( ) <EOL> class BasePwizTestCase ( PeeweeTestCase ) : <EOL> models = [ ] <EOL> def setUp ( self ) : <EOL> super ( BasePwizTestCase , self ) . setUp ( ) <EOL> if os . path . exists ( db . database ) : <EOL> os . unlink ( db . database ) <EOL> db . connect ( ) <EOL> db . create_tables ( self . models ) <EOL> self . introspector = Introspector . from_database ( db ) <EOL> def tearDown ( self ) : <EOL> super ( BasePwizTestCase , self ) . tearDown ( ) <EOL> db . close ( ) <EOL> class TestPwiz ( BasePwizTestCase ) : <EOL> models = [ User , Note , Category ] <EOL> def test_print_models ( self ) : <EOL> with capture_output ( ) as output : <EOL> print_models ( self . introspector ) <EOL> self . assertEqual ( output . data . strip ( ) , EXPECTED ) <EOL> def test_print_header ( self ) : <EOL> cmdline = '<STR_LIT>' % db . database <EOL> with capture_output ( ) as output : <EOL> with mock . patch ( '<STR_LIT>' ) as mock_datetime : <EOL> now = mock_datetime . now . return_value <EOL> now . strftime . return_value = '<STR_LIT>' <EOL> print_header ( cmdline , self . introspector ) <EOL> self . assertEqual ( output . data . strip ( ) , ( <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' ) % ( cmdline , db . database , peewee_version ) ) <EOL> @ skip_if ( lambda : sys . version_info [ : <NUM_LIT:2> ] < ( <NUM_LIT:2> , <NUM_LIT:7> ) ) <EOL> class TestPwizOrdered ( BasePwizTestCase ) : <EOL> models = [ User , Note ] <EOL> def test_ordered_columns ( self ) : <EOL> with capture_output ( ) as output : <EOL> print_models ( self . introspector , preserve_order = True ) <EOL> self . assertEqual ( output . data . strip ( ) , EXPECTED_ORDERED ) </s>
