Merge pull request #8 from erik1110/main

[feat] add timeout
tomy0000000 · Aug 21, 2023 · 496e14d · 496e14d
2 parents eae1428 + 3f3c29b
commit 496e14d
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 9 deletions.
diff --git a/tests/test_app_api_client.py b/tests/test_app_api_client.py
@@ -21,6 +21,7 @@
 TEST_TIME = 1655654400
 TEST_TS_TOLERANCE = 30
 TEST_UUID = "test_uuid"
+TEST_TIMEOUT = (3, 1)
 
 
 @pytest.fixture
@@ -86,6 +87,7 @@ def test_get_lottery_numbers(client, mocker):
             "UUID": TEST_UUID,
             "appID": TEST_APP_ID,
         },
+        timeout=TEST_TIMEOUT,
     )
     mocked_check_api_error.assert_called_once()
     mocked_parse_obj.assert_called_once()
@@ -135,6 +137,7 @@ def test_get_invoice_header(client, mocker):
                 "UUID": TEST_UUID,
                 "appID": TEST_APP_ID,
             },
+            timeout=TEST_TIMEOUT,
         )
         mocked_check_api_error.assert_called_once()
         mocked_parse_obj.assert_called_once()
@@ -251,6 +254,7 @@ def test_get_invoice_detail(client, mocker):
             "randomNumber": TEST_INVOICE_RANDOM,
             "appID": TEST_APP_ID,
         },
+        timeout=TEST_TIMEOUT,
     )
     mocked_check_api_error.assert_called_once()
     mocked_parse_obj.assert_called_once()
@@ -285,6 +289,7 @@ def test_get_invoice_detail(client, mocker):
             "randomNumber": TEST_INVOICE_RANDOM,
             "appID": TEST_APP_ID,
         },
+        timeout=TEST_TIMEOUT,
     )
     mocked_check_api_error.assert_called_once()
     mocked_parse_obj.assert_called_once()
@@ -306,6 +311,7 @@ def test_get_love_code(client, mocker):
             "UUID": TEST_UUID,
             "appID": TEST_APP_ID,
         },
+        timeout=TEST_TIMEOUT,
     )
     mocked_check_api_error.assert_called_once()
     mocked_parse_obj.assert_called_once()
@@ -344,6 +350,7 @@ def test_get_carrier_invoices_header(client, mocker):
             "appID": TEST_APP_ID,
             "cardEncrypt": TEST_CARD_ENCRYPT,
         },
+        timeout=TEST_TIMEOUT,
     )
     mocked_check_api_error.assert_called_once()
     mocked_parse_obj.assert_called_once()
@@ -393,6 +400,7 @@ def test_get_carrier_invoices_detail(client, mocker):
             "appID": TEST_APP_ID,
             "cardEncrypt": TEST_CARD_ENCRYPT,
         },
+        timeout=TEST_TIMEOUT,
     )
     mocked_check_api_error.assert_called_once()
     mocked_parse_obj.assert_called_once()
@@ -444,6 +452,7 @@ def test_carrier_donate_invoice(client, mocker):
             "cardEncrypt": TEST_CARD_ENCRYPT,
             "signature": "Ydmt4alIcCLYG/VBcI+QUdyfgO5Yp8bfX6TPy5xgsVs=",
         },
+        timeout=TEST_TIMEOUT,
     )
     mocked_check_api_error.assert_called_once()
     mocked_parse_obj.assert_called_once()
@@ -479,6 +488,7 @@ def test_get_aggregate_carrier(client, mocker):
             "uuid": TEST_UUID,
             "signature": "2pWN1GfP6S7oncE56OJetvpxGlE1tFYHeqgNwHmIuw4=",
         },
+        timeout=TEST_TIMEOUT,
     )
     mocked_check_api_error.assert_called_once()
     mocked_parse_obj.assert_called_once()

diff --git a/tw_invoice/app_client.py b/tw_invoice/app_client.py
@@ -1,6 +1,6 @@
 from datetime import date
 from time import time
-from typing import Union
+from typing import Tuple, Union
 from uuid import uuid4
 
 try:
@@ -40,6 +40,7 @@ def __init__(
         ts_tolerance: int = 20,
         max_retries: int = 20,
         skip_validation: bool = False,
+        timeout: Union[float, Tuple[float, float], Tuple[float, None]] = (3, 1),
     ):
         self.app_id = app_id
         self.api_key = api_key
@@ -67,6 +68,7 @@ def __init__(
                 )
             ),
         )
+        self.timeout = timeout
 
     def get_lottery_numbers(
         self, invoice_term: str
@@ -83,7 +85,9 @@ def get_lottery_numbers(
             "UUID": self.uuid,
             "appID": self.app_id,
         }
-        results = check_api_error(self.session.post(URL, data=data))
+        results = check_api_error(
+            self.session.post(URL, data=data, timeout=self.timeout)
+        )
         if not self.skip_validation:
             results = LotteryNumberResponse.parse_obj(results)
         return results
@@ -111,7 +115,9 @@ def get_invoice_header(
             "UUID": self.uuid,
             "appID": self.app_id,
         }
-        results = check_api_error(self.session.post(URL, data=data))
+        results = check_api_error(
+            self.session.post(URL, data=data, timeout=self.timeout)
+        )
         if not self.skip_validation:
             results = InvoiceHeaderResponse.parse_obj(results)
         return results
@@ -183,7 +189,9 @@ def get_invoice_detail(
             "randomNumber": invoice_random,
             "appID": self.app_id,
         }
-        results = check_api_error(self.session.post(URL, data=data))
+        results = check_api_error(
+            self.session.post(URL, data=data, timeout=self.timeout)
+        )
         if not self.skip_validation:
             results = InvoiceDetailResponse.parse_obj(results)
         return results
@@ -199,7 +207,9 @@ def get_love_code(self, query: str) -> dict:
             "UUID": self.uuid,
             "appID": self.app_id,
         }
-        results = check_api_error(self.session.post(URL, data=data))
+        results = check_api_error(
+            self.session.post(URL, data=data, timeout=self.timeout)
+        )
         if not self.skip_validation:
             results = LoveCodeResponse.parse_obj(results)
         return results
@@ -230,7 +240,9 @@ def get_carrier_invoices_header(
             "appID": self.app_id,
             "cardEncrypt": card_encrypt,
         }
-        results = check_api_error(self.session.post(URL, data=data))
+        results = check_api_error(
+            self.session.post(URL, data=data, timeout=self.timeout)
+        )
         if not self.skip_validation:
             results = CarrierInvoicesHeaderResponse.parse_obj(results)
         return results
@@ -265,7 +277,9 @@ def get_carrier_invoices_detail(
             "appID": self.app_id,
             "cardEncrypt": card_encrypt,
         }
-        results = check_api_error(self.session.post(URL, data=data))
+        results = check_api_error(
+            self.session.post(URL, data=data, timeout=self.timeout)
+        )
         if not self.skip_validation:
             results = CarrierInvoicesDetailResponse.parse_obj(results)
         return results
@@ -302,7 +316,9 @@ def carrier_donate_invoice(
         signature = sign(data, self.api_key)
         data["signature"] = signature
         self.serial += 1
-        results = check_api_error(self.session.post(URL, data=data))
+        results = check_api_error(
+            self.session.post(URL, data=data, timeout=self.timeout)
+        )
         if not self.skip_validation:
             results = CarrierInvoiceDonateResponse.parse_obj(results)
         return results
@@ -330,7 +346,9 @@ def get_aggregate_carrier(
         signature = sign(data, self.api_key)
         data["signature"] = signature
         self.serial += 1
-        results = check_api_error(self.session.post(URL, data=data))
+        results = check_api_error(
+            self.session.post(URL, data=data, timeout=self.timeout)
+        )
         if not self.skip_validation:
             results = AggregateCarrierResponse.parse_obj(results)
         return results