<s> import re <EOL> from django . conf import settings <EOL> from django . db import connection <EOL> from . compat import get_model_name <EOL> class assert_number_queries ( object ) : <EOL> def __init__ ( self , number ) : <EOL> self . number = number <EOL> def matched_queries ( self ) : <EOL> return connection . queries <EOL> def query_count ( self ) : <EOL> return len ( self . matched_queries ( ) ) <EOL> def __enter__ ( self ) : <EOL> self . DEBUG = settings . DEBUG <EOL> settings . DEBUG = True <EOL> self . num_queries_before = self . query_count ( ) <EOL> def __exit__ ( self , type , value , traceback ) : <EOL> self . num_queries_after = self . query_count ( ) <EOL> assert self . num_queries_after - self . num_queries_before == self . number <EOL> settings . DEBUG = self . DEBUG <EOL> class RegexMixin ( object ) : <EOL> regex = None <EOL> def matched_queries ( self ) : <EOL> matched_queries = super ( RegexMixin , self ) . matched_queries ( ) <EOL> if self . regex is not None : <EOL> pattern = re . compile ( self . regex ) <EOL> regex_compliant_queries = [ query for query in matched_queries if pattern . match ( query . get ( '<STR_LIT>' ) ) ] <EOL> return regex_compliant_queries <EOL> class assert_number_of_queries_on_regex ( RegexMixin , assert_number_queries ) : <EOL> def __init__ ( self , regex , number ) : <EOL> super ( assert_number_of_queries_on_regex , self ) . __init__ ( number ) <EOL> self . regex = regex <EOL> class assert_select_number_queries_on_model ( assert_number_of_queries_on_regex ) : <EOL> def __init__ ( self , model_class , number ) : <EOL> model_name = get_model_name ( model_class ) <EOL> regex = r'<STR_LIT>' % model_name <EOL> super ( assert_select_number_queries_on_model , self ) . __init__ ( regex , number ) </s>
