import time
import unittest

from mock import Mock
from trollius import coroutine, Return, get_event_loop, From

from util.metrics.metricqueue import duration_collector_async


mock_histogram = Mock()

class NonReturn(Exception):
  pass 


@coroutine
@duration_collector_async(mock_histogram, labelvalues=["testlabel"])
def duration_decorated():
    time.sleep(1)
    raise Return("fin")


@coroutine
@duration_collector_async(mock_histogram, labelvalues=["testlabel"])
def duration_decorated_error():
    raise NonReturn("not a Return error")

@coroutine
def calls_decorated():
  yield From(duration_decorated())


class DurationDecoratorTestCase(unittest.TestCase):
  def __init__(self, *args, **kwargs):
    self.etcd_client_mock = None
    self.loop = get_event_loop()
    super(DurationDecoratorTestCase, self).__init__(*args, **kwargs)

  def test_duration_decorator(self):
    self.loop.run_until_complete(duration_decorated())
    assert mock_histogram.Observe.called
    assert 1 - mock_histogram.Observe.call_args[0][0] < 1  # duration should be close to 1s
    assert mock_histogram.Observe.call_args[1]["labelvalues"] == ["testlabel"]
    
  def test_duration_decorator_error(self):
    mock_histogram.reset_mock()
    
    with self.assertRaises(NonReturn):
      self.loop.run_until_complete(duration_decorated_error())
    assert not mock_histogram.Observe.called
    
  def test_duration_decorator_caller(self):
    mock_histogram.reset_mock()
    
    self.loop.run_until_complete(calls_decorated())
    assert mock_histogram.Observe.called
    assert 1 - mock_histogram.Observe.call_args[0][0] < 1  # duration should be close to 1s
    assert mock_histogram.Observe.call_args[1]["labelvalues"] == ["testlabel"]
    

if __name__ == '__main__':
  unittest.main()