Merge pull request #97 from kannkyo/feature-reduce-wait-time

Reduce wait time

SECURITY.md

@@ -12,7 +12,7 @@ currently being supported with security updates.
 ## Reporting a Vulnerability
 
 
-Please report a vulunerability by [security advisories](https://github.com/kannkyo/nvd-api/security/advisories/new). Vulunerability reports should contain 4 elements.
+Please report a vulnerability by [security advisories](https://github.com/kannkyo/nvd-api/security/advisories/new). Vulunerability reports should contain 4 elements.
 
 * Impact : What kind of vulnerability is it? Who is impacted?
 * Patches : Has the problem been patched? What versions should users upgrade to?

src/nvd_api/client.py

@@ -73,7 +73,7 @@ def __init__(self,
         """Constructor  # noqa: E501
 
         Args:
-            wait_time (int, optional): wait time (ms) after api execution. Defaults to 6,000.
+            wait_time (int, optional): max wait time (ms) after api execution. Defaults to 6,000.
             request_timeout (int, optional): API request time (ms). Defaults to 20,000.
             max_retries (int, optional): max API retry count. Defaults to 3.
             api_key (str, optional): NVD API 2.0 key
@@ -89,6 +89,7 @@ def __init__(self,
         self._products_api = ProductsApi(api_client)
         self.wait_time = wait_time
         self.request_timeout = request_timeout
+        self._last_exec_datetime = datetime.now()
 
     def get_cves(self,
                  cpe_name: str = None,
@@ -218,9 +219,10 @@ def get_cves(self,
         kwargs = {k: v for k, v in kwargs.items() if v is not None}  # Noneは削除
 
         logger.debug(f"execute cves api : kwargs={kwargs}")
+        self._sleep()
         ret = self._vulnerabilities_api.get_cves(**kwargs)
+        self.update_last_exec_datetime()
         logger.debug(f"execute cves api : response={ret}")
-        self._sleep()
 
         return ret
 
@@ -263,9 +265,10 @@ def get_cve_history(self,
         kwargs = {k: v for k, v in kwargs.items() if v is not None}  # Noneは削除
 
         logger.debug(f"execute cve history api : kwargs={kwargs}")
+        self._sleep()
         ret = self._vulnerabilities_api.get_cve_history(**kwargs)
+        self.update_last_exec_datetime()
         logger.debug(f"execute cve history api : response={ret}")
-        self._sleep()
 
         return ret
 
@@ -317,9 +320,10 @@ def get_cpes(self,
         kwargs = {k: v for k, v in kwargs.items() if v is not None}  # Noneは削除
 
         logger.debug(f"execute cpes api : kwargs={kwargs}")
+        self._sleep()
         ret = self._products_api.get_cpes(**kwargs)
+        self.update_last_exec_datetime()
         logger.debug(f"execute cpes api : response={ret}")
-        self._sleep()
 
         return ret
 
@@ -359,9 +363,10 @@ def get_cpe_match(self,
         kwargs = {k: v for k, v in kwargs.items() if v is not None}  # Noneは削除
 
         logger.debug(f"execute cpe match api : kwargs={kwargs}")
+        self._sleep()
         ret = self._products_api.get_cpe_match(**kwargs)
+        self.update_last_exec_datetime()
         logger.debug(f"execute cpe match api : response={ret}")
-        self._sleep()
 
         return ret
 
@@ -724,6 +729,9 @@ def _verify_keyword(self,
 
     def _sleep(self, wait_time: int = None):
         if wait_time is None:
-            time.sleep(self.wait_time/1000)
-        else:
-            time.sleep(wait_time/1000)
+            wait_time = self.wait_time
+        dt = datetime.now() - self._last_exec_datetime
+        time.sleep(max(0, wait_time/1000 - dt.total_seconds()))
+
+    def update_last_exec_datetime(self):
+        self._last_exec_datetime = datetime.now()

tests/test_sleep.py

@@ -9,20 +9,37 @@
 """
 
 
+import time
 import unittest
 
 from nvd_api.client import NvdApiClient
 
 
 class TestGetCpes(unittest.TestCase):
     def setUp(self):
-        self.client = NvdApiClient()  # noqa: E501
+        pass
 
     def tearDown(self):
         pass
 
-    def test_sleep(self):
-        self.client._sleep(10)
+    def test_sleep_with_wait_time(self):
+        client = NvdApiClient(wait_time=1000)
+        start = time.time()
+        client._sleep(1000)
+        end = time.time()
+
+        duration = end - start
+        assert duration >= 1
+
+    def test_sleep_without_wait_time(self):
+        client = NvdApiClient(wait_time=1000)
+        time.sleep(1.1)
+        start = time.time()
+        client._sleep()
+        end = time.time()
+
+        duration = end - start
+        assert duration <= 0.1
 
 
 if __name__ == '__main__':