def test_dynamic_get(self):
    entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                              "prop2val", ns="blah")
    zookeeper = flexmock()
    zookeeper.should_receive("get_valid_transaction_id").and_return(1)
    zookeeper.should_receive("register_updated_key").and_return(1)
    zookeeper.should_receive("acquire_lock").and_return(True)
    db_batch = flexmock()
    db_batch.should_receive('valid_data_version').and_return(True)
    db_batch.should_receive("batch_put_entity").and_return(None)
    db_batch.should_receive("batch_get_entity").and_return(
               {"test\x00blah\x00test_kind:nancy\x01": 
                 {
                   APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                   APP_ENTITY_SCHEMA[1]: 1
                 }
               })
    db_batch.should_receive('record_reads')

    dd = DatastoreDistributed(db_batch, zookeeper) 

    entity_key = entity_proto1.key()
    get_req = datastore_pb.GetRequest()
    key = get_req.add_key() 
    key.MergeFrom(entity_key)
    get_resp = datastore_pb.GetResponse()
    
    dd.dynamic_get("test", get_req, get_resp)     
    self.assertEquals(get_resp.entity_size(), 1)

    # Now test while in a transaction
    get_resp = datastore_pb.GetResponse()
    get_req.mutable_transaction().set_handle(1)
    dd.dynamic_get("test", get_req, get_resp)     
    self.assertEquals(get_resp.entity_size(), 1)
  def test_dynamic_get(self):
    entity_proto1 = self.get_new_entity_proto("test", "test_kind", "nancy", "prop1name", 
                                              "prop2val", ns="blah")
    zookeeper = flexmock()
    zookeeper.should_receive("get_valid_transaction_id").and_return(1)
    zookeeper.should_receive("register_updated_key").and_return(1)
    zookeeper.should_receive("acquire_lock").and_return(True)
    db_batch = flexmock()
    db_batch.should_receive('valid_data_version').and_return(True)
    db_batch.should_receive("batch_put_entity").and_return(None)
    db_batch.should_receive("batch_get_entity").and_return(
               {"test\x00blah\x00test_kind:nancy\x01": 
                 {
                   APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
                   APP_ENTITY_SCHEMA[1]: 1
                 }
               })

    dd = DatastoreDistributed(db_batch, zookeeper) 

    entity_key = entity_proto1.key()
    get_req = datastore_pb.GetRequest()
    key = get_req.add_key() 
    key.MergeFrom(entity_key)
    get_resp = datastore_pb.GetResponse()
    
    dd.dynamic_get("test", get_req, get_resp)     
    self.assertEquals(get_resp.entity_size(), 1)

    # Now test while in a transaction
    get_resp = datastore_pb.GetResponse()
    get_req.mutable_transaction().set_handle(1)
    dd.dynamic_get("test", get_req, get_resp)     
    self.assertEquals(get_resp.entity_size(), 1)
  def test_dynamic_get(self):
    entity_proto1 = self.get_new_entity_proto(
      "test", "test_kind", "nancy", "prop1name", "prop2val", ns="blah")

    zk_client = flexmock()
    zk_client.should_receive('add_listener')

    async_result = gen.Future()
    async_result.set_result({
      "test\x00blah\x00test_kind:nancy\x01": {
        APP_ENTITY_SCHEMA[0]: entity_proto1.Encode(),
        APP_ENTITY_SCHEMA[1]: 1
      }
    })

    zookeeper = flexmock(handle=zk_client)
    zookeeper.should_receive("get_valid_transaction_id").and_return(1)
    zookeeper.should_receive("acquire_lock").and_return(True)
    db_batch = flexmock()
    db_batch.should_receive('valid_data_version_sync').and_return(True)
    db_batch.should_receive("batch_get_entity").and_return(async_result)
    db_batch.should_receive('record_reads').and_return(ASYNC_NONE)

    transaction_manager = flexmock()
    dd = DatastoreDistributed(db_batch, transaction_manager, zookeeper)

    entity_key = entity_proto1.key()
    get_req = datastore_pb.GetRequest()
    key = get_req.add_key()
    key.MergeFrom(entity_key)
    get_resp = datastore_pb.GetResponse()

    yield dd.dynamic_get("test", get_req, get_resp)
    self.assertEquals(get_resp.entity_size(), 1)

    # Now test while in a transaction
    get_resp = datastore_pb.GetResponse()
    get_req.mutable_transaction().set_handle(1)
    yield dd.dynamic_get("test", get_req, get_resp)
    self.assertEquals(get_resp.entity_size(), 1)