def test_acquire_locks_for_trans(self): zk_client = flexmock() zk_client.should_receive('add_listener') zookeeper = flexmock(handle=zk_client) db_batch = flexmock() db_batch.should_receive('valid_data_version_sync').and_return(True) transaction_manager = flexmock() dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper) flexmock(dd).should_receive("is_instance_wrapper").and_return(False).once() self.assertRaises(TypeError, dd.acquire_locks_for_trans, [1], 1) dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper) flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \ .and_return(False).and_return(False) self.assertRaises(TypeError, dd.acquire_locks_for_trans, [1], 1) dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper) flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \ .and_return(True) dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper) flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \ .and_return(True).and_return(False) flexmock(dd).should_receive("get_table_prefix").and_return("prefix").never() flexmock(dd).should_receive("get_root_key_from_entity_key").and_return("rootkey").never() self.assertEquals({}, dd.acquire_locks_for_trans([], 1)) zookeeper = flexmock(handle=zk_client) zookeeper.should_receive("acquire_lock").once() dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper) entity = flexmock() entity.should_receive("app").and_return("appid") flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \ .and_return(True).and_return(True) flexmock(dd).should_receive("get_root_key_from_entity_key").and_return("rootkey").once() self.assertEquals({'rootkey':1}, dd.acquire_locks_for_trans([entity], 1)) zookeeper = flexmock(handle=zk_client) zookeeper.should_receive("acquire_lock").once().and_raise(ZKTransactionException) zookeeper.should_receive("notify_failed_transaction").once() dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper) entity = flexmock() entity.should_receive("app").and_return("appid") flexmock(dd).should_receive("is_instance_wrapper").and_return(True) \ .and_return(True).and_return(True) flexmock(dd).should_receive("get_root_key_from_entity_key").and_return("rootkey").once() self.assertRaises(ZKTransactionException, dd.acquire_locks_for_trans, [entity], 1)