def test_context(self):
try:
with statesampler.instruction_id('A'):
tracker = statesampler.for_test()
with tracker.scoped_state(NameContext('name', 'tid'), 'stage'):
_LOGGER.info('message a')
with statesampler.instruction_id('B'):
_LOGGER.info('message b')
_LOGGER.info('message c')
self.fn_log_handler.close()
a, b, c = sum([
list(logs.log_entries)
for logs in self.test_logging_service.log_records_received
], [])
self.assertEqual(a.instruction_id, 'A')
self.assertEqual(b.instruction_id, 'B')
self.assertEqual(c.instruction_id, '')
self.assertEqual(a.transform_id, 'tid')
self.assertEqual(b.transform_id, '')
self.assertEqual(c.transform_id, '')
finally:
statesampler.set_current_tracker(None)
def test_metrics(self):
sampler = statesampler.StateSampler('', counters.CounterFactory())
statesampler.set_current_tracker(sampler)
state1 = sampler.scoped_state(
'mystep', 'myState', metrics_container=MetricsContainer('mystep'))
try:
sampler.start()
with state1:
counter = MetricTests.base_metric_group.counter("my_counter")
meter = MetricTests.base_metric_group.meter("my_meter")
distribution = MetricTests.base_metric_group.distribution("my_distribution")
container = MetricsEnvironment.current_container()
self.assertEqual(0, counter.get_count())
self.assertEqual(0, meter.get_count())
self.assertEqual(
DistributionData(
0, 0, 0, 0), container.get_distribution(
MetricName(
'[]', 'my_distribution')).get_cumulative())
counter.inc(-2)
meter.mark_event(3)
distribution.update(10)
distribution.update(2)
self.assertEqual(-2, counter.get_count())
self.assertEqual(3, meter.get_count())
self.assertEqual(
DistributionData(
12, 2, 2, 10), container.get_distribution(
MetricName(
'[]', 'my_distribution')).get_cumulative())
finally:
sampler.stop()
def test_create_counter_distribution(self):
sampler = statesampler.StateSampler('', counters.CounterFactory())
statesampler.set_current_tracker(sampler)
state1 = sampler.scoped_state('mystep', 'myState',
metrics_container=MetricsContainer('mystep'))
sampler.start()
with state1:
counter_ns = 'aCounterNamespace'
distro_ns = 'aDistributionNamespace'
name = 'a_name'
counter = Metrics.counter(counter_ns, name)
distro = Metrics.distribution(distro_ns, name)
counter.inc(10)
counter.dec(3)
distro.update(10)
distro.update(2)
self.assertTrue(isinstance(counter, Metrics.DelegatingCounter))
self.assertTrue(isinstance(distro, Metrics.DelegatingDistribution))
del distro
del counter
container = MetricsEnvironment.current_container()
self.assertEqual(
container.counters[MetricName(counter_ns, name)].get_cumulative(),
7)
self.assertEqual(
container.distributions[MetricName(distro_ns, name)].get_cumulative(),
DistributionData(12, 2, 2, 10))
sampler.stop()
def test_create_process_wide(self):
sampler = statesampler.StateSampler('', counters.CounterFactory())
statesampler.set_current_tracker(sampler)
state1 = sampler.scoped_state(
'mystep', 'myState', metrics_container=MetricsContainer('mystep'))
try:
sampler.start()
with state1:
urn = "my:custom:urn"
labels = {'key': 'value'}
counter = InternalMetrics.counter(
urn=urn, labels=labels, process_wide=True)
# Test that if process_wide is set, that it will be set
# on the process_wide container.
counter.inc(10)
self.assertTrue(isinstance(counter, Metrics.DelegatingCounter))
del counter
metric_name = MetricName(None, None, urn=urn, labels=labels)
# Expect a value set on the current container.
self.assertEqual(
MetricsEnvironment.process_wide_container().get_counter(
metric_name).get_cumulative(),
10)
# Expect no value set on the current container.
self.assertEqual(
MetricsEnvironment.current_container().get_counter(
metric_name).get_cumulative(),
0)
finally:
sampler.stop()
def test_nested_with_per_thread_info(self):
self.maxDiff = None
tracker = statesampler.StateSampler('stage', CounterFactory())
statesampler.set_current_tracker(tracker)
formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid')
with logger.PerThreadLoggingContext(work_item_id='workitem'):
with tracker.scoped_state('step1', 'process'):
record = self.create_log_record(**self.SAMPLE_RECORD)
log_output1 = json.loads(formatter.format(record))
with tracker.scoped_state('step2', 'process'):
record = self.create_log_record(**self.SAMPLE_RECORD)
log_output2 = json.loads(formatter.format(record))
record = self.create_log_record(**self.SAMPLE_RECORD)
log_output3 = json.loads(formatter.format(record))
statesampler.set_current_tracker(None)
record = self.create_log_record(**self.SAMPLE_RECORD)
log_output4 = json.loads(formatter.format(record))
self.assertEqual(
log_output1,
dict(self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step1'))
self.assertEqual(
log_output2,
dict(self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step2'))
self.assertEqual(
log_output3,
dict(self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step1'))
self.assertEqual(log_output4, self.SAMPLE_OUTPUT)
def test_nested_with_per_thread_info(self):
self.maxDiff = None
tracker = statesampler.StateSampler('stage', CounterFactory())
statesampler.set_current_tracker(tracker)
formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid')
with logger.PerThreadLoggingContext(work_item_id='workitem'):
with tracker.scoped_state('step1', 'process'):
record = self.create_log_record(**self.SAMPLE_RECORD)
log_output1 = json.loads(formatter.format(record))
with tracker.scoped_state('step2', 'process'):
record = self.create_log_record(**self.SAMPLE_RECORD)
log_output2 = json.loads(formatter.format(record))
record = self.create_log_record(**self.SAMPLE_RECORD)
log_output3 = json.loads(formatter.format(record))
statesampler.set_current_tracker(None)
record = self.create_log_record(**self.SAMPLE_RECORD)
log_output4 = json.loads(formatter.format(record))
self.assertEqual(log_output1, dict(
self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step1'))
self.assertEqual(log_output2, dict(
self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step2'))
self.assertEqual(log_output3, dict(
self.SAMPLE_OUTPUT, work='workitem', stage='stage', step='step1'))
self.assertEqual(log_output4, self.SAMPLE_OUTPUT)
def test_create_counter_distribution(self):
sampler = statesampler.StateSampler('', counters.CounterFactory())
statesampler.set_current_tracker(sampler)
state1 = sampler.scoped_state('mystep', 'myState',
metrics_container=MetricsContainer('mystep'))
sampler.start()
with state1:
counter_ns = 'aCounterNamespace'
distro_ns = 'aDistributionNamespace'
name = 'a_name'
counter = Metrics.counter(counter_ns, name)
distro = Metrics.distribution(distro_ns, name)
counter.inc(10)
counter.dec(3)
distro.update(10)
distro.update(2)
self.assertTrue(isinstance(counter, Metrics.DelegatingCounter))
self.assertTrue(isinstance(distro, Metrics.DelegatingDistribution))
del distro
del counter
container = MetricsEnvironment.current_container()
self.assertEqual(
container.counters[MetricName(counter_ns, name)].get_cumulative(),
7)
self.assertEqual(
container.distributions[MetricName(distro_ns, name)].get_cumulative(),
DistributionData(12, 2, 2, 10))
sampler.stop()
def run(self):
state_sampler = statesampler.StateSampler('', counters.CounterFactory())
statesampler.set_current_tracker(state_sampler)
while not self.shutdown_requested:
task = self._get_task_or_none()
if task:
try:
if not self.shutdown_requested:
self._update_name(task)
task.call(state_sampler)
self._update_name()
finally:
self.queue.task_done()
def run(self):
state_sampler = statesampler.StateSampler('', counters.CounterFactory())
statesampler.set_current_tracker(state_sampler)
while not self.shutdown_requested:
task = self._get_task_or_none()
if task:
try:
if not self.shutdown_requested:
self._update_name(task)
task.call(state_sampler)
self._update_name()
finally:
self.queue.task_done()
def test_record_with_per_thread_info(self):
self.maxDiff = None
tracker = statesampler.StateSampler('stage', CounterFactory())
statesampler.set_current_tracker(tracker)
formatter = logger.JsonLogFormatter(job_id='jobid', worker_id='workerid')
with logger.PerThreadLoggingContext(work_item_id='workitem'):
with tracker.scoped_state('step', 'process'):
record = self.create_log_record(**self.SAMPLE_RECORD)
log_output = json.loads(formatter.format(record))
expected_output = dict(self.SAMPLE_OUTPUT)
expected_output.update(
{'work': 'workitem', 'stage': 'stage', 'step': 'step'})
self.assertEqual(log_output, expected_output)
statesampler.set_current_tracker(None)
