<s> import datetime <EOL> import os <EOL> import random <EOL> import string <EOL> import time <EOL> import unittest <EOL> from google . appengine . ext import db <EOL> from mapreduce import control <EOL> from mapreduce import hooks <EOL> from mapreduce import model <EOL> from mapreduce import test_support <EOL> from testlib import testutil <EOL> from mapreduce . api import map_job <EOL> def random_string ( length ) : <EOL> """<STR_LIT>""" <EOL> return "<STR_LIT>" . join ( <EOL> random . choice ( string . letters + string . digits ) for _ in range ( length ) ) <EOL> class TestEntity ( db . Model ) : <EOL> """<STR_LIT>""" <EOL> class TestHooks ( hooks . Hooks ) : <EOL> """<STR_LIT>""" <EOL> enqueue_kickoff_task_calls = [ ] <EOL> def enqueue_kickoff_task ( self , task , queue_name ) : <EOL> TestHooks . enqueue_kickoff_task_calls . append ( ( task , queue_name ) ) <EOL> def test_handler ( entity ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> class ControlTest ( testutil . HandlerTestBase ) : <EOL> """<STR_LIT>""" <EOL> QUEUE_NAME = "<STR_LIT>" <EOL> def setUp ( self ) : <EOL> testutil . HandlerTestBase . setUp ( self ) <EOL> TestHooks . enqueue_kickoff_task_calls = [ ] <EOL> def get_mapreduce_spec ( self , task ) : <EOL> """<STR_LIT>""" <EOL> payload = test_support . decode_task_payload ( task ) <EOL> return model . MapreduceSpec . from_json_str ( payload [ "<STR_LIT>" ] ) <EOL> def validate_map_started ( self , mapreduce_id , queue_name = None ) : <EOL> """<STR_LIT>""" <EOL> queue_name = queue_name or self . QUEUE_NAME <EOL> self . assertTrue ( mapreduce_id ) <EOL> tasks = self . taskqueue . GetTasks ( queue_name ) <EOL> self . assertEquals ( <NUM_LIT:1> , len ( tasks ) ) <EOL> task = tasks [ <NUM_LIT:0> ] <EOL> self . assertEqual ( "<STR_LIT>" + mapreduce_id , <EOL> task [ "<STR_LIT:url>" ] ) <EOL> handler = test_support . execute_task ( task ) <EOL> self . assertEqual ( mapreduce_id , handler . request . get ( "<STR_LIT>" ) ) <EOL> state = model . MapreduceState . get_by_job_id ( mapreduce_id ) <EOL> params = map_job . JobConfig . _get_default_mr_params ( ) <EOL> params . update ( { "<STR_LIT:foo>" : "<STR_LIT:bar>" , <EOL> "<STR_LIT>" : "<STR_LIT>" , <EOL> "<STR_LIT>" : queue_name } ) <EOL> self . assertEqual ( state . mapreduce_spec . params , params ) <EOL> job_config = map_job . JobConfig . _to_map_job_config ( state . mapreduce_spec , <EOL> queue_name ) <EOL> self . assertEqual ( <NUM_LIT:0> , job_config . _api_version ) <EOL> return task [ "<STR_LIT>" ] <EOL> def testStartMap ( self ) : <EOL> """<STR_LIT>""" <EOL> TestEntity ( ) . put ( ) <EOL> shard_count = <NUM_LIT:4> <EOL> mapreduce_id = control . start_map ( <EOL> "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> { <EOL> "<STR_LIT>" : __name__ + "<STR_LIT:.>" + TestEntity . __name__ , <EOL> } , <EOL> shard_count , <EOL> mapreduce_parameters = { "<STR_LIT:foo>" : "<STR_LIT:bar>" } , <EOL> base_path = "<STR_LIT>" , <EOL> queue_name = self . QUEUE_NAME ) <EOL> self . validate_map_started ( mapreduce_id ) <EOL> def testStartMap_Countdown ( self ) : <EOL> """<STR_LIT>""" <EOL> TestEntity ( ) . put ( ) <EOL> now_sec = long ( time . time ( ) ) <EOL> shard_count = <NUM_LIT:4> <EOL> mapreduce_id = control . start_map ( <EOL> "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> { <EOL> "<STR_LIT>" : __name__ + "<STR_LIT:.>" + TestEntity . __name__ , <EOL> } , <EOL> shard_count , <EOL> mapreduce_parameters = { "<STR_LIT:foo>" : "<STR_LIT:bar>" } , <EOL> base_path = "<STR_LIT>" , <EOL> queue_name = self . QUEUE_NAME , <EOL> countdown = <NUM_LIT:1000> ) <EOL> task_eta = self . validate_map_started ( mapreduce_id ) <EOL> eta_sec = time . mktime ( time . strptime ( task_eta , "<STR_LIT>" ) ) <EOL> self . assertTrue ( now_sec + <NUM_LIT:1000> <= eta_sec ) <EOL> def testStartMap_Eta ( self ) : <EOL> """<STR_LIT>""" <EOL> TestEntity ( ) . put ( ) <EOL> eta = datetime . datetime . utcnow ( ) + datetime . timedelta ( hours = <NUM_LIT:1> ) <EOL> shard_count = <NUM_LIT:4> <EOL> mapreduce_id = control . start_map ( <EOL> "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> { <EOL> "<STR_LIT>" : __name__ + "<STR_LIT:.>" + TestEntity . __name__ , <EOL> } , <EOL> shard_count , <EOL> mapreduce_parameters = { "<STR_LIT:foo>" : "<STR_LIT:bar>" } , <EOL> base_path = "<STR_LIT>" , <EOL> queue_name = self . QUEUE_NAME , <EOL> eta = eta ) <EOL> task_eta = self . validate_map_started ( mapreduce_id ) <EOL> self . assertEquals ( eta . strftime ( "<STR_LIT>" ) , task_eta ) <EOL> def testStartMap_QueueEnvironment ( self ) : <EOL> """<STR_LIT>""" <EOL> TestEntity ( ) . put ( ) <EOL> shard_count = <NUM_LIT:4> <EOL> os . environ [ "<STR_LIT>" ] = self . QUEUE_NAME <EOL> try : <EOL> mapreduce_id = control . start_map ( <EOL> "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> { <EOL> "<STR_LIT>" : __name__ + "<STR_LIT:.>" + TestEntity . __name__ , <EOL> } , <EOL> shard_count , <EOL> mapreduce_parameters = { "<STR_LIT:foo>" : "<STR_LIT:bar>" } , <EOL> base_path = "<STR_LIT>" ) <EOL> finally : <EOL> del os . environ [ "<STR_LIT>" ] <EOL> self . validate_map_started ( mapreduce_id ) <EOL> def testStartMap_Hooks ( self ) : <EOL> """<STR_LIT>""" <EOL> TestEntity ( ) . put ( ) <EOL> shard_count = <NUM_LIT:4> <EOL> mapreduce_id = control . start_map ( <EOL> "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> { <EOL> "<STR_LIT>" : __name__ + "<STR_LIT:.>" + TestEntity . __name__ , <EOL> } , <EOL> shard_count , <EOL> mapreduce_parameters = { "<STR_LIT:foo>" : "<STR_LIT:bar>" } , <EOL> base_path = "<STR_LIT>" , <EOL> queue_name = "<STR_LIT>" , <EOL> hooks_class_name = __name__ + "<STR_LIT:.>" + TestHooks . __name__ ) <EOL> self . assertTrue ( mapreduce_id ) <EOL> task , queue_name = TestHooks . enqueue_kickoff_task_calls [ <NUM_LIT:0> ] <EOL> self . assertEqual ( "<STR_LIT>" + mapreduce_id , <EOL> task . url ) <EOL> self . assertEqual ( "<STR_LIT>" , queue_name ) <EOL> def testStartMap_RaisingHooks ( self ) : <EOL> """<STR_LIT>""" <EOL> TestEntity ( ) . put ( ) <EOL> shard_count = <NUM_LIT:4> <EOL> mapreduce_id = control . start_map ( <EOL> "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> { <EOL> "<STR_LIT>" : __name__ + "<STR_LIT:.>" + TestEntity . __name__ , <EOL> } , <EOL> shard_count , <EOL> mapreduce_parameters = { "<STR_LIT:foo>" : "<STR_LIT:bar>" } , <EOL> base_path = "<STR_LIT>" , <EOL> queue_name = "<STR_LIT>" , <EOL> hooks_class_name = hooks . __name__ + "<STR_LIT:.>" + hooks . Hooks . __name__ ) <EOL> self . validate_map_started ( mapreduce_id ) <EOL> def testStartMap_HugePayload ( self ) : <EOL> """<STR_LIT>""" <EOL> TestEntity ( ) . put ( ) <EOL> shard_count = <NUM_LIT:4> <EOL> mapreduce_id = "<STR_LIT>" <EOL> mapreduce_id = control . start_map ( <EOL> "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> { <EOL> "<STR_LIT>" : __name__ + "<STR_LIT:.>" + TestEntity . __name__ , <EOL> "<STR_LIT>" : random_string ( <NUM_LIT> ) <EOL> } , <EOL> shard_count , <EOL> mapreduce_parameters = { "<STR_LIT:foo>" : "<STR_LIT:bar>" } , <EOL> base_path = "<STR_LIT>" , <EOL> queue_name = self . QUEUE_NAME ) <EOL> self . validate_map_started ( mapreduce_id ) <EOL> def testStartMapTransactional ( self ) : <EOL> """<STR_LIT>""" <EOL> TestEntity ( ) . put ( ) <EOL> shard_count = <NUM_LIT:4> <EOL> mapreduce_id = "<STR_LIT>" <EOL> @ db . transactional ( xg = True ) <EOL> def tx ( ) : <EOL> some_entity = TestEntity ( ) <EOL> some_entity . put ( ) <EOL> return control . start_map ( <EOL> "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> { <EOL> "<STR_LIT>" : __name__ + "<STR_LIT:.>" + TestEntity . __name__ , <EOL> } , <EOL> shard_count , <EOL> mapreduce_parameters = { "<STR_LIT:foo>" : "<STR_LIT:bar>" } , <EOL> base_path = "<STR_LIT>" , <EOL> queue_name = self . QUEUE_NAME , <EOL> in_xg_transaction = True ) <EOL> mapreduce_id = tx ( ) <EOL> self . validate_map_started ( mapreduce_id ) <EOL> def testStartMapTransactional_HugePayload ( self ) : <EOL> """<STR_LIT>""" <EOL> TestEntity ( ) . put ( ) <EOL> shard_count = <NUM_LIT:4> <EOL> mapreduce_id = "<STR_LIT>" <EOL> @ db . transactional ( xg = True ) <EOL> def tx ( ) : <EOL> some_entity = TestEntity ( ) <EOL> some_entity . put ( ) <EOL> return control . start_map ( <EOL> "<STR_LIT>" , <EOL> __name__ + "<STR_LIT>" , <EOL> "<STR_LIT>" , <EOL> { <EOL> "<STR_LIT>" : __name__ + "<STR_LIT:.>" + TestEntity . __name__ , <EOL> "<STR_LIT>" : random_string ( <NUM_LIT> ) <EOL> } , <EOL> shard_count , <EOL> mapreduce_parameters = { "<STR_LIT:foo>" : "<STR_LIT:bar>" } , <EOL> base_path = "<STR_LIT>" , <EOL> queue_name = self . QUEUE_NAME , <EOL> in_xg_transaction = True ) <EOL> mapreduce_id = tx ( ) <EOL> self . validate_map_started ( mapreduce_id ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> unittest . main ( ) </s>
