<s> import sys <EOL> sys . path [ <NUM_LIT:0> : <NUM_LIT:0> ] = [ "<STR_LIT>" ] <EOL> import unittest <EOL> from mongoengine import * <EOL> from mongoengine . connection import get_db <EOL> from mongoengine . context_managers import ( switch_db , switch_collection , <EOL> no_sub_classes , no_dereference , <EOL> query_counter ) <EOL> class ContextManagersTest ( unittest . TestCase ) : <EOL> def test_switch_db_context_manager ( self ) : <EOL> connect ( '<STR_LIT>' ) <EOL> register_connection ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> class Group ( Document ) : <EOL> name = StringField ( ) <EOL> Group . drop_collection ( ) <EOL> Group ( name = "<STR_LIT>" ) . save ( ) <EOL> self . assertEqual ( <NUM_LIT:1> , Group . objects . count ( ) ) <EOL> with switch_db ( Group , '<STR_LIT>' ) as Group : <EOL> self . assertEqual ( <NUM_LIT:0> , Group . objects . count ( ) ) <EOL> Group ( name = "<STR_LIT:hello>" ) . save ( ) <EOL> self . assertEqual ( <NUM_LIT:1> , Group . objects . count ( ) ) <EOL> Group . drop_collection ( ) <EOL> self . assertEqual ( <NUM_LIT:0> , Group . objects . count ( ) ) <EOL> self . assertEqual ( <NUM_LIT:1> , Group . objects . count ( ) ) <EOL> def test_switch_collection_context_manager ( self ) : <EOL> connect ( '<STR_LIT>' ) <EOL> register_connection ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> class Group ( Document ) : <EOL> name = StringField ( ) <EOL> Group . drop_collection ( ) <EOL> with switch_collection ( Group , '<STR_LIT>' ) as Group : <EOL> Group . drop_collection ( ) <EOL> Group ( name = "<STR_LIT>" ) . save ( ) <EOL> self . assertEqual ( <NUM_LIT:1> , Group . objects . count ( ) ) <EOL> with switch_collection ( Group , '<STR_LIT>' ) as Group : <EOL> self . assertEqual ( <NUM_LIT:0> , Group . objects . count ( ) ) <EOL> Group ( name = "<STR_LIT>" ) . save ( ) <EOL> self . assertEqual ( <NUM_LIT:1> , Group . objects . count ( ) ) <EOL> Group . drop_collection ( ) <EOL> self . assertEqual ( <NUM_LIT:0> , Group . objects . count ( ) ) <EOL> self . assertEqual ( <NUM_LIT:1> , Group . objects . count ( ) ) <EOL> def test_no_dereference_context_manager_object_id ( self ) : <EOL> """<STR_LIT>""" <EOL> connect ( '<STR_LIT>' ) <EOL> class User ( Document ) : <EOL> name = StringField ( ) <EOL> class Group ( Document ) : <EOL> ref = ReferenceField ( User , dbref = False ) <EOL> generic = GenericReferenceField ( ) <EOL> members = ListField ( ReferenceField ( User , dbref = False ) ) <EOL> User . drop_collection ( ) <EOL> Group . drop_collection ( ) <EOL> for i in xrange ( <NUM_LIT:1> , <NUM_LIT> ) : <EOL> User ( name = '<STR_LIT>' % i ) . save ( ) <EOL> user = User . objects . first ( ) <EOL> Group ( ref = user , members = User . objects , generic = user ) . save ( ) <EOL> with no_dereference ( Group ) as NoDeRefGroup : <EOL> self . assertTrue ( Group . _fields [ '<STR_LIT>' ] . _auto_dereference ) <EOL> self . assertFalse ( NoDeRefGroup . _fields [ '<STR_LIT>' ] . _auto_dereference ) <EOL> with no_dereference ( Group ) as Group : <EOL> group = Group . objects . first ( ) <EOL> self . assertTrue ( all ( [ not isinstance ( m , User ) <EOL> for m in group . members ] ) ) <EOL> self . assertFalse ( isinstance ( group . ref , User ) ) <EOL> self . assertFalse ( isinstance ( group . generic , User ) ) <EOL> self . assertTrue ( all ( [ isinstance ( m , User ) <EOL> for m in group . members ] ) ) <EOL> self . assertTrue ( isinstance ( group . ref , User ) ) <EOL> self . assertTrue ( isinstance ( group . generic , User ) ) <EOL> def test_no_dereference_context_manager_dbref ( self ) : <EOL> """<STR_LIT>""" <EOL> connect ( '<STR_LIT>' ) <EOL> class User ( Document ) : <EOL> name = StringField ( ) <EOL> class Group ( Document ) : <EOL> ref = ReferenceField ( User , dbref = True ) <EOL> generic = GenericReferenceField ( ) <EOL> members = ListField ( ReferenceField ( User , dbref = True ) ) <EOL> User . drop_collection ( ) <EOL> Group . drop_collection ( ) <EOL> for i in xrange ( <NUM_LIT:1> , <NUM_LIT> ) : <EOL> User ( name = '<STR_LIT>' % i ) . save ( ) <EOL> user = User . objects . first ( ) <EOL> Group ( ref = user , members = User . objects , generic = user ) . save ( ) <EOL> with no_dereference ( Group ) as NoDeRefGroup : <EOL> self . assertTrue ( Group . _fields [ '<STR_LIT>' ] . _auto_dereference ) <EOL> self . assertFalse ( NoDeRefGroup . _fields [ '<STR_LIT>' ] . _auto_dereference ) <EOL> with no_dereference ( Group ) as Group : <EOL> group = Group . objects . first ( ) <EOL> self . assertTrue ( all ( [ not isinstance ( m , User ) <EOL> for m in group . members ] ) ) <EOL> self . assertFalse ( isinstance ( group . ref , User ) ) <EOL> self . assertFalse ( isinstance ( group . generic , User ) ) <EOL> self . assertTrue ( all ( [ isinstance ( m , User ) <EOL> for m in group . members ] ) ) <EOL> self . assertTrue ( isinstance ( group . ref , User ) ) <EOL> self . assertTrue ( isinstance ( group . generic , User ) ) <EOL> def test_no_sub_classes ( self ) : <EOL> class A ( Document ) : <EOL> x = IntField ( ) <EOL> y = IntField ( ) <EOL> meta = { '<STR_LIT>' : True } <EOL> class B ( A ) : <EOL> z = IntField ( ) <EOL> class C ( B ) : <EOL> zz = IntField ( ) <EOL> A . drop_collection ( ) <EOL> A ( x = <NUM_LIT:10> , y = <NUM_LIT:20> ) . save ( ) <EOL> A ( x = <NUM_LIT:15> , y = <NUM_LIT:30> ) . save ( ) <EOL> B ( x = <NUM_LIT:20> , y = <NUM_LIT> ) . save ( ) <EOL> B ( x = <NUM_LIT:30> , y = <NUM_LIT:50> ) . save ( ) <EOL> C ( x = <NUM_LIT> , y = <NUM_LIT> ) . save ( ) <EOL> self . assertEqual ( A . objects . count ( ) , <NUM_LIT:5> ) <EOL> self . assertEqual ( B . objects . count ( ) , <NUM_LIT:3> ) <EOL> self . assertEqual ( C . objects . count ( ) , <NUM_LIT:1> ) <EOL> with no_sub_classes ( A ) as A : <EOL> self . assertEqual ( A . objects . count ( ) , <NUM_LIT:2> ) <EOL> for obj in A . objects : <EOL> self . assertEqual ( obj . __class__ , A ) <EOL> with no_sub_classes ( B ) as B : <EOL> self . assertEqual ( B . objects . count ( ) , <NUM_LIT:2> ) <EOL> for obj in B . objects : <EOL> self . assertEqual ( obj . __class__ , B ) <EOL> with no_sub_classes ( C ) as C : <EOL> self . assertEqual ( C . objects . count ( ) , <NUM_LIT:1> ) <EOL> for obj in C . objects : <EOL> self . assertEqual ( obj . __class__ , C ) <EOL> self . assertEqual ( A . objects . count ( ) , <NUM_LIT:5> ) <EOL> self . assertEqual ( B . objects . count ( ) , <NUM_LIT:3> ) <EOL> self . assertEqual ( C . objects . count ( ) , <NUM_LIT:1> ) <EOL> def test_query_counter ( self ) : <EOL> connect ( '<STR_LIT>' ) <EOL> db = get_db ( ) <EOL> db . test . find ( { } ) <EOL> with query_counter ( ) as q : <EOL> self . assertEqual ( <NUM_LIT:0> , q ) <EOL> for i in xrange ( <NUM_LIT:1> , <NUM_LIT> ) : <EOL> db . test . find ( { } ) . count ( ) <EOL> self . assertEqual ( <NUM_LIT:50> , q ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> unittest . main ( ) </s>
