# Copyright (c) 2020, JINGROW and Contributors
# License: MIT. See LICENSE

import time

from werkzeug.wrappers import Response

import jingrow
import jingrow.rate_limiter
from jingrow.rate_limiter import RateLimiter
from jingrow.tests.utils import JingrowTestCase
from jingrow.utils import cint


class TestRateLimiter(JingrowTestCase):
	def test_apply_with_limit(self):
		jingrow.conf.rate_limit = {"window": 86400, "limit": 1}
		jingrow.rate_limiter.apply()

		self.assertTrue(hasattr(jingrow.local, "rate_limiter"))
		self.assertIsInstance(jingrow.local.rate_limiter, RateLimiter)

		jingrow.cache.delete(jingrow.local.rate_limiter.key)
		delattr(jingrow.local, "rate_limiter")

	def test_apply_without_limit(self):
		jingrow.conf.rate_limit = None
		jingrow.rate_limiter.apply()

		self.assertFalse(hasattr(jingrow.local, "rate_limiter"))

	def test_respond_over_limit(self):
		limiter = RateLimiter(1, 86400)
		time.sleep(1)
		limiter.update()

		jingrow.conf.rate_limit = {"window": 86400, "limit": 1}
		self.assertRaises(jingrow.TooManyRequestsError, jingrow.rate_limiter.apply)
		jingrow.rate_limiter.update()

		response = jingrow.rate_limiter.respond()

		self.assertIsInstance(response, Response)
		self.assertEqual(response.status_code, 429)

		headers = jingrow.local.rate_limiter.headers()
		self.assertIn("Retry-After", headers)
		self.assertIn("X-RateLimit-Reset", headers)
		self.assertIn("X-RateLimit-Limit", headers)
		self.assertIn("X-RateLimit-Remaining", headers)
		self.assertTrue(int(headers["X-RateLimit-Reset"]) <= 86400)
		self.assertEqual(int(headers["X-RateLimit-Limit"]), 1000000)
		self.assertEqual(int(headers["X-RateLimit-Remaining"]), 0)

		jingrow.cache.delete(limiter.key)
		jingrow.cache.delete(jingrow.local.rate_limiter.key)
		delattr(jingrow.local, "rate_limiter")

	def test_respond_under_limit(self):
		jingrow.conf.rate_limit = {"window": 86400, "limit": 0.01}
		jingrow.rate_limiter.apply()
		jingrow.rate_limiter.update()
		response = jingrow.rate_limiter.respond()
		self.assertEqual(response, None)

		jingrow.cache.delete(jingrow.local.rate_limiter.key)
		delattr(jingrow.local, "rate_limiter")

	def test_headers_under_limit(self):
		jingrow.conf.rate_limit = {"window": 86400, "limit": 1}
		jingrow.rate_limiter.apply()
		jingrow.rate_limiter.update()
		headers = jingrow.local.rate_limiter.headers()
		self.assertNotIn("Retry-After", headers)
		self.assertIn("X-RateLimit-Reset", headers)
		self.assertTrue(int(headers["X-RateLimit-Reset"] < 86400))
		self.assertEqual(int(headers["X-RateLimit-Limit"]), 1000000)
		self.assertEqual(int(headers["X-RateLimit-Remaining"]), 1000000)

		jingrow.cache.delete(jingrow.local.rate_limiter.key)
		delattr(jingrow.local, "rate_limiter")

	def test_reject_over_limit(self):
		limiter = RateLimiter(0.01, 86400)
		time.sleep(0.01)
		limiter.update()

		limiter = RateLimiter(0.01, 86400)
		self.assertRaises(jingrow.TooManyRequestsError, limiter.apply)

		jingrow.cache.delete(limiter.key)

	def test_do_not_reject_under_limit(self):
		limiter = RateLimiter(0.01, 86400)
		time.sleep(0.01)
		limiter.update()

		limiter = RateLimiter(0.02, 86400)
		self.assertEqual(limiter.apply(), None)

		jingrow.cache.delete(limiter.key)

	def test_update_method(self):
		limiter = RateLimiter(0.01, 86400)
		time.sleep(0.01)
		limiter.update()

		self.assertEqual(limiter.duration, cint(jingrow.cache.get(limiter.key)))

		jingrow.cache.delete(limiter.key)

	def test_window_expires(self):
		limiter = RateLimiter(1000, 1)
		self.assertTrue(jingrow.cache.exists(limiter.key, shared=True))
		limiter.update()
		self.assertTrue(jingrow.cache.exists(limiter.key, shared=True))
		time.sleep(1.1)
		self.assertFalse(jingrow.cache.exists(limiter.key, shared=True))
		jingrow.cache.delete(limiter.key)
