def check_estimated_size_bytes(self, entity_bytes, timestamp, namespace=None): """A helper method to test get_estimated_size_bytes""" timestamp_req = helper.make_request( self._PROJECT, namespace, helper.make_latest_timestamp_query(namespace)) timestamp_resp = self.make_stats_response( {'timestamp': datastore_helper.from_timestamp(timestamp)}) kind_stat_req = helper.make_request( self._PROJECT, namespace, helper.make_kind_stats_query( namespace, self._query.kind[0].name, datastore_helper.micros_from_timestamp(timestamp))) kind_stat_resp = self.make_stats_response( {'entity_bytes': entity_bytes}) def fake_run_query(req): if req == timestamp_req: return timestamp_resp elif req == kind_stat_req: return kind_stat_resp else: print kind_stat_req raise ValueError("Unknown req: %s" % req) self._mock_datastore.run_query.side_effect = fake_run_query self.assertEqual(entity_bytes, ReadFromDatastore.get_estimated_size_bytes( self._PROJECT, namespace, self._query, self._mock_datastore)) self.assertEqual(self._mock_datastore.run_query.call_args_list, [call(timestamp_req), call(kind_stat_req)])