def test_is_identified(self):
with LogDir() as log_dir:
KM.log_dir = log_dir
with StdIO() as stdio:
self.assertFalse(KM.is_identified())
KM.identify('id')
self.assertTrue(KM.is_identified())
def test_generate_query(self):
host = 'example.com'
_t = 1
_k = 'key'
_p = 'id'
with LogDir() as log_dir:
KM.init(_k, host=host, to_stderr=False, log_dir=log_dir,
use_cron=False)
KM.identify(_p)
# Ensure time.time() always returns the same value
time = self.mocker.replace('time.time')
time()
self.mocker.result(_t)
self.mocker.count(1, None)
urlopen = self.mocker.replace('urllib2.urlopen')
# First
urlopen('http://%s/e?_t=%d&_k=%s' % (host, _t, _k))
self.mocker.result(None)
# Second
urlopen('http://%s/e?_t=%d&_k=%s&_p=%s' % (host, _t, _k, _p))
self.mocker.result(None)
# Third
urlopen('http://%s/f?_t=%d&_k=%s&_p=%s' % (host, _t, _k, _p))
self.mocker.throw(urllib2.HTTPError('', 500, 'Error', None, None))
with self.mocker:
# First
KM.generate_query('e', {}, update=False)
# Second
KM.generate_query('e', {}, update=True)
# Third
KM.generate_query('f', {}, update=True)
with self.get_log('error') as log:
self.assertEndsWith(log.read(),
'> HTTP Error 500: Error\n')
def test_is_identified(self):
with LogDir() as log_dir:
KM.log_dir = log_dir
with StdIO() as stdio:
self.assertFalse(KM.is_identified())
KM.identify('id')
self.assertTrue(KM.is_identified())
def test_is_initialized(self):
with LogDir() as log_dir:
KM.log_dir = log_dir
with StdIO() as stdio:
self.assertFalse(KM.is_initialized())
KM.init('key')
self.assertTrue(KM.is_initialized())
def test_record(self):
action = 'action'
with self.assertQuery('e', {'_n': action}):
KM.record(action)
data = {'a': 'a'}
with self.assertQuery('s', data):
KM.record(data)
def test_record(self):
action = 'action'
with self.assertQuery('e', {'_n': action}):
KM.record(action)
data = {'a': 'a'}
with self.assertQuery('s', data):
KM.record(data)
def test_is_initialized(self):
with LogDir() as log_dir:
KM.log_dir = log_dir
with StdIO() as stdio:
self.assertFalse(KM.is_initialized())
KM.init('key')
self.assertTrue(KM.is_initialized())
def assertQuery(self, type, data, update=True):
host = 'example.com'
_t = 1
_k = 'key'
_p = 'id'
with LogDir() as log_dir:
KM.init(_k,
host=host,
to_stderr=False,
log_dir=log_dir,
use_cron=True)
KM.identify(_p)
# Ensure time.time() always returns the same value
time = self.mocker.replace('time.time')
time()
self.mocker.result(_t)
self.mocker.count(0, None)
with self.mocker:
yield
with self.get_log('query') as log:
data = dict(data, _t=str(_t), _k=_k)
if update:
data['_p'] = _p
line = log.readline().strip()
bits = urlparse.urlsplit(line)
self.assertEqual(bits[2], '/%s' % urllib.quote(type))
self.assertEqual(dict(parse_qsl(bits[3])), data)
self.assertEqual(log.readline(), '')
os.unlink(KM.log_name('query'))
def assertQuery(self, type, data, update=True):
host = 'example.com'
_t = 1
_k = 'key'
_p = 'id'
with LogDir() as log_dir:
KM.init(_k, host=host, to_stderr=False, log_dir=log_dir,
use_cron=True)
KM.identify(_p)
# Ensure time.time() always returns the same value
time = self.mocker.replace('time.time')
time()
self.mocker.result(_t)
self.mocker.count(0, None)
with self.mocker:
yield
with self.get_log('query') as log:
data = dict(data, _t=str(_t), _k=_k)
if update:
data['_p'] = _p
line = log.readline().strip()
bits = urlparse.urlsplit(line)
self.assertEqual(bits[2], '/%s' % urllib.quote(type))
self.assertEqual(dict(parse_qsl(bits[3])), data)
self.assertEqual(log.readline(), '')
os.unlink(KM.log_name('query'))
def test_log_dir_writable(self):
log_dir = '/invalid/'
with StdIO() as stdio:
KM.init('key', to_stderr=True, log_dir=log_dir)
self.assertEqual(stdio.stderr.getvalue(),
("Couldn't open %(log_dir)skissmetrics_query.log "
"for writing. Does %(log_dir)s exist? "
"Permissions?\n" % {'log_dir': log_dir}))
def test_log_error(self):
with LogDir() as log_dir:
KM.init('key', to_stderr=True, log_dir=log_dir)
with StdIO() as stdio:
KM.log_error('Error')
self.assertEndsWith(stdio.stderr.getvalue(), '> Error\n')
with self.get_log('error') as log:
self.assertEndsWith(log.read(), '> Error\n')
def test_log_error(self):
with LogDir() as log_dir:
KM.init('key', to_stderr=True, log_dir=log_dir)
with StdIO() as stdio:
KM.log_error('Error')
self.assertEndsWith(stdio.stderr.getvalue(),
'> Error\n')
with self.get_log('error') as log:
self.assertEndsWith(log.read(), '> Error\n')
def test_log_dir_writable(self):
log_dir = '/invalid/'
with StdIO() as stdio:
KM.init('key', to_stderr=True, log_dir=log_dir)
self.assertEqual(stdio.stderr.getvalue(),
("Couldn't open %(log_dir)skissmetrics_query.log "
"for writing. Does %(log_dir)s exist? "
"Permissions?\n" % {
'log_dir': log_dir
}))
def test_is_initialize_and_identified(self):
with LogDir() as log_dir:
KM.log_dir = log_dir
with StdIO() as stdio:
self.assertFalse(KM.is_initialized_and_identified())
KM.init('key')
self.assertFalse(KM.is_initialized_and_identified())
KM.identify('id')
self.assertTrue(KM.is_initialized_and_identified())
KM.init(None)
self.assertFalse(KM.is_initialized_and_identified())
def test_is_initialize_and_identified(self):
with LogDir() as log_dir:
KM.log_dir = log_dir
with StdIO() as stdio:
self.assertFalse(KM.is_initialized_and_identified())
KM.init('key')
self.assertFalse(KM.is_initialized_and_identified())
KM.identify('id')
self.assertTrue(KM.is_initialized_and_identified())
KM.init(None)
self.assertFalse(KM.is_initialized_and_identified())
def test_log(self):
with LogDir() as log_dir:
KM.init('key', to_stderr=True, log_dir=log_dir)
KM.log('error', 'Error')
with self.get_log('error') as log:
self.assertEndsWith(log.read(), 'Error\n')
KM.log('query', 'Query')
with self.get_log('query') as log:
self.assertEndsWith(log.read(), 'Query\n')
KM.log('send', 'Send')
with self.get_log('send') as log:
self.assertEndsWith(log.read(), 'Send\n')
KM.log('invalid', 'Invalid')
def test_log(self):
with LogDir() as log_dir:
KM.init('key', to_stderr=True, log_dir=log_dir)
KM.log('error', 'Error')
with self.get_log('error') as log:
self.assertEndsWith(log.read(), 'Error\n')
KM.log('query', 'Query')
with self.get_log('query') as log:
self.assertEndsWith(log.read(), 'Query\n')
KM.log('send', 'Send')
with self.get_log('send') as log:
self.assertEndsWith(log.read(), 'Send\n')
KM.log('invalid', 'Invalid')
def test_generate_query_cron(self):
host = 'example.com'
_t = 1
_k = 'key'
_p = 'id'
with LogDir() as log_dir:
KM.init(_k,
host=host,
to_stderr=False,
log_dir=log_dir,
use_cron=True)
KM.identify(_p)
# Ensure time.time() always returns the same value
time = self.mocker.replace('time.time')
time()
self.mocker.result(_t)
self.mocker.count(1, None)
with self.mocker:
KM.generate_query('e', {}, update=False)
KM.generate_query('e', {}, update=True)
with self.get_log('query') as log:
self.assertEqual(log.readline(),
'/e?_t=%s&_k=%s\n' % (_t, _k))
self.assertEqual(log.readline(),
'/e?_t=%s&_k=%s&_p=%s\n' % (_t, _k, _p))
self.assertEqual(log.read(), '')
def get_log(self, log_name):
try:
with open(KM.log_name(log_name), 'r') as log_file:
yield log_file
except IOError, e:
if e.errno == errno.ENOENT:
yield None
def get_log(self, log_name):
try:
with open(KM.log_name(log_name), 'r') as log_file:
yield log_file
except IOError, e:
if e.errno == errno.ENOENT:
yield None
def test_send_query(self):
host = 'example.com'
line = '/line?param=1'
url = 'http://%s%s' % (host, line)
with LogDir() as log_dir:
KM.init('key', host=host, to_stderr=False, log_dir=log_dir)
urlopen = self.mocker.replace('urllib2.urlopen')
# First
urlopen(url)
self.mocker.result(None)
# Second
urlopen(url)
self.mocker.throw(urllib2.HTTPError(url, 500, 'Error', None, None))
with self.mocker:
# First
KM.send_query(line)
# Second
self.assertRaises(urllib2.HTTPError, KM.send_query, line)
def test_send_query(self):
host = 'example.com'
line = '/line?param=1'
url = 'http://%s%s' % (host, line)
with LogDir() as log_dir:
KM.init('key', host=host, to_stderr=False, log_dir=log_dir)
urlopen = self.mocker.replace('urllib2.urlopen')
# First
urlopen(url)
self.mocker.result(None)
# Second
urlopen(url)
self.mocker.throw(urllib2.HTTPError(url, 500, 'Error', None, None))
with self.mocker:
# First
KM.send_query(line)
# Second
self.assertRaises(urllib2.HTTPError, KM.send_query, line)
def dispatch(self, dispatch_observ):
""" Compute the assignment between drivers and passengers at each time step
:param dispatch_observ: a list of dict, the key in the dict includes:
order_id, int
driver_id, int
order_driver_distance, float
order_start_location, a list as [lng, lat], float
order_finish_location, a list as [lng, lat], float
driver_location, a list as [lng, lat], float
timestamp, int
order_finish_timestamp, int
day_of_week, int
reward_units, float
pick_up_eta, float
:return: a list of dict, the key in the dict includes:
order_id and driver_id, the pair indicating the assignment
"""
if len(dispatch_observ) == 0:
return []
order_id2loc, driver_id2loc = {}, {}
order_loc2id, driver_loc2id = {}, {}
count_order, count_driver = 0, 0
for pair in dispatch_observ:
#if self.get_finish_prob(pair) <= 0: # TODO: may eliminate an order/driver?
# continue
if pair['order_id'] not in order_id2loc:
order_id2loc[pair['order_id']] = count_order
order_loc2id[count_order] = pair['order_id']
count_order += 1
if pair['driver_id'] not in driver_id2loc:
driver_id2loc[pair['driver_id']] = count_driver
driver_loc2id[count_driver] = pair['driver_id']
count_driver += 1
weights = torch.zeros(len(order_id2loc),
len(driver_id2loc) + len(order_id2loc))
for pair in dispatch_observ:
i = order_id2loc[pair['order_id']]
j = driver_id2loc[pair['driver_id']]
weights[i, j] = self.get_weight(pair)
raw_assign = KM(weights).run()[:len(driver_id2loc)]
dispatch_action = []
for i, item in enumerate(raw_assign):
if item != -1:
dispatch_action.append({
'order_id': order_loc2id[item.item()],
'driver_id': driver_loc2id[i]
})
self.update_value_network(dispatch_observ, dispatch_action)
return dispatch_action
def test_reset(self):
with LogDir() as log_dir:
KM.init('key', to_stderr=False, log_dir=log_dir)
KM.identify('id')
KM.log_query('Query')
self.assertEqual(KM._key, 'key')
self.assertEqual(KM._id, 'id')
self.assertEqual(
KM._logs,
{'query': os.path.join(log_dir, 'kissmetrics_query.log')})
KM.reset()
self.assertEqual(KM._key, None)
self.assertEqual(KM._id, None)
self.assertEqual(KM._logs, {})
def test_send_logged_queries(self):
host = 'example.com'
line = '/type?foo=1&bar=2'
with LogDir() as log_dir:
KM.init('key', host=host, log_dir=log_dir)
# Write line to the query log
KM.log_query(line)
# Expect that line to be requested
urlopen = self.mocker.replace('urllib2.urlopen')
urlopen('http://%s%s' % (host, line))
self.mocker.result(None)
with self.mocker:
KM.send_logged_queries()
self.assertFalse(os.path.exists(KM.log_name('query')))
self.assertFalse(os.path.exists(KM.log_name('send')))
def test_reset(self):
with LogDir() as log_dir:
KM.init('key', to_stderr=False, log_dir=log_dir)
KM.identify('id')
KM.log_query('Query')
self.assertEqual(KM._key, 'key')
self.assertEqual(KM._id, 'id')
self.assertEqual(
KM._logs,
{'query': os.path.join(log_dir, 'kissmetrics_query.log')}
)
KM.reset()
self.assertEqual(KM._key, None)
self.assertEqual(KM._id, None)
self.assertEqual(KM._logs, {})
def test_send_logged_queries(self):
host = 'example.com'
line = '/type?foo=1&bar=2'
with LogDir() as log_dir:
KM.init('key', host=host, log_dir=log_dir)
# Write line to the query log
KM.log_query(line)
# Expect that line to be requested
urlopen = self.mocker.replace('urllib2.urlopen')
urlopen('http://%s%s' % (host, line))
self.mocker.result(None)
with self.mocker:
KM.send_logged_queries()
self.assertFalse(os.path.exists(KM.log_name('query')))
self.assertFalse(os.path.exists(KM.log_name('send')))
def test_log_name(self):
log_dir = '/var/log'
KM.init('key', to_stderr=False, log_dir=log_dir)
self.assertEqual(KM.log_name('invalid'), os.path.join(log_dir, ''))
self.assertEqual(KM.log_name('error'),
os.path.join(log_dir, 'kissmetrics_error.log'))
self.assertEqual(KM.log_name('query'),
os.path.join(log_dir, 'kissmetrics_query.log'))
send_name = KM.log_name('send')
self.assertStartsWith(send_name, log_dir)
self.assertEndsWith(send_name, 'kissmetrics_sending.log')
self.assertEqual(KM.log_name('send'), send_name)
def test_generate_query(self):
host = 'example.com'
_t = 1
_k = 'key'
_p = 'id'
with LogDir() as log_dir:
KM.init(_k,
host=host,
to_stderr=False,
log_dir=log_dir,
use_cron=False)
KM.identify(_p)
# Ensure time.time() always returns the same value
time = self.mocker.replace('time.time')
time()
self.mocker.result(_t)
self.mocker.count(1, None)
urlopen = self.mocker.replace('urllib2.urlopen')
# First
urlopen('http://%s/e?_t=%d&_k=%s' % (host, _t, _k))
self.mocker.result(None)
# Second
urlopen('http://%s/e?_t=%d&_k=%s&_p=%s' % (host, _t, _k, _p))
self.mocker.result(None)
# Third
urlopen('http://%s/f?_t=%d&_k=%s&_p=%s' % (host, _t, _k, _p))
self.mocker.throw(urllib2.HTTPError('', 500, 'Error', None, None))
with self.mocker:
# First
KM.generate_query('e', {}, update=False)
# Second
KM.generate_query('e', {}, update=True)
# Third
KM.generate_query('f', {}, update=True)
with self.get_log('error') as log:
self.assertEndsWith(log.read(),
'> HTTP Error 500: Error\n')
def test_log_name(self):
log_dir = '/var/log'
KM.init('key', to_stderr=False, log_dir=log_dir)
self.assertEqual(KM.log_name('invalid'),
os.path.join(log_dir, ''))
self.assertEqual(KM.log_name('error'),
os.path.join(log_dir, 'kissmetrics_error.log'))
self.assertEqual(KM.log_name('query'),
os.path.join(log_dir, 'kissmetrics_query.log'))
send_name = KM.log_name('send')
self.assertStartsWith(send_name, log_dir)
self.assertEndsWith(send_name, 'kissmetrics_sending.log')
self.assertEqual(KM.log_name('send'), send_name)
def test_generate_query_cron(self):
host = 'example.com'
_t = 1
_k = 'key'
_p = 'id'
with LogDir() as log_dir:
KM.init(_k, host=host, to_stderr=False, log_dir=log_dir,
use_cron=True)
KM.identify(_p)
# Ensure time.time() always returns the same value
time = self.mocker.replace('time.time')
time()
self.mocker.result(_t)
self.mocker.count(1, None)
with self.mocker:
KM.generate_query('e', {}, update=False)
KM.generate_query('e', {}, update=True)
with self.get_log('query') as log:
self.assertEqual(log.readline(),
'/e?_t=%s&_k=%s\n' % (_t, _k))
self.assertEqual(log.readline(),
'/e?_t=%s&_k=%s&_p=%s\n' % (_t, _k, _p))
self.assertEqual(log.read(), '')
def test_set(self):
data = {'a': 'a'}
with self.assertQuery('s', data):
KM.set(data)
def test_log_query(self):
with LogDir() as log_dir:
KM.init('key', to_stderr=False, log_dir=log_dir)
KM.log_query('Query')
with self.get_log('query') as log:
self.assertEndsWith(log.read(), 'Query\n')
def test_init(self):
self.assertRaises(TypeError, KM.init)
KM.init('key1')
self.assertEqual(KM._key, 'key1')
KM.init(key='key2')
self.assertEqual(KM._key, 'key2')
KM.init('key', host='example.com:80')
self.assertEqual(KM.host, 'example.com:80')
KM.init('key', log_dir='/tmp/')
self.assertEqual(KM.log_dir, '/tmp/')
KM.init('key', use_cron=True)
self.assertEqual(KM._use_cron, True)
KM.init('key', to_stderr=False)
self.assertEqual(KM._to_stderr, False)
def test_log_send(self):
with LogDir() as log_dir:
KM.init('key', to_stderr=False, log_dir=log_dir)
KM.log_send('Send')
with self.get_log('send') as log:
self.assertEndsWith(log.read(), 'Send\n')
def test_log_query(self):
with LogDir() as log_dir:
KM.init('key', to_stderr=False, log_dir=log_dir)
KM.log_query('Query')
with self.get_log('query') as log:
self.assertEndsWith(log.read(), 'Query\n')
def test_set(self):
data = {'a': 'a'}
with self.assertQuery('s', data):
KM.set(data)
def test_alias(self):
name = 'name'
alias_to = 'alias_to'
with self.assertQuery('a', {'_n': alias_to, '_p': name}, update=False):
KM.alias(name, alias_to)
def test_log_send(self):
with LogDir() as log_dir:
KM.init('key', to_stderr=False, log_dir=log_dir)
KM.log_send('Send')
with self.get_log('send') as log:
self.assertEndsWith(log.read(), 'Send\n')
def test_init(self):
self.assertRaises(TypeError, KM.init)
KM.init('key1')
self.assertEqual(KM._key, 'key1')
KM.init(key='key2')
self.assertEqual(KM._key, 'key2')
KM.init('key', host='example.com:80')
self.assertEqual(KM.host, 'example.com:80')
KM.init('key', log_dir='/tmp/')
self.assertEqual(KM.log_dir, '/tmp/')
KM.init('key', use_cron=True)
self.assertEqual(KM._use_cron, True)
KM.init('key', to_stderr=False)
self.assertEqual(KM._to_stderr, False)
def test_alias(self):
name = 'name'
alias_to = 'alias_to'
with self.assertQuery('a', {'_n': alias_to, '_p': name}, update=False):
KM.alias(name, alias_to)
def test_log_dir(self):
KM.init('key', to_stderr=False, log_dir='/invalid/')
def test_log_dir(self):
KM.init('key', to_stderr=False, log_dir='/invalid/')