get files as bytes (unit-finance#239)

axshani · ilyamerman · web-flow · commit 9cf7d5ce56b9 · 2023-12-30T23:40:16.000+02:00
Co-authored-by: Ilya Merman &lt;ilya.merman@gmail.com&gt;
diff --git a/e2e_tests/statement_test.py b/e2e_tests/statement_test.py
@@ -6,25 +6,24 @@
 token = os.environ.get('TOKEN')
 client = Unit("https://api.s.unit.sh", token)
 
+
 def test_list_and_get_statements():
     statements = client.statements.list(ListStatementParams(2)).data
     for s in statements:
         assert s.type == "accountStatementDTO"
 
         params = GetStatementParams(s.id)
         html_statement = client.statements.get(params).data
-        assert "<!DOCTYPE html>" in html_statement
+        assert "<!DOCTYPE html>" in str(html_statement)
 
         params = GetStatementParams(s.id, customer_id=s.relationships["customer"].id)
         html_statement = client.statements.get(params).data
-        assert "<!DOCTYPE html>" in html_statement
+        assert "<!DOCTYPE html>" in str(html_statement)
 
         account_id = s.relationships["account"].id
         pdf_response = client.statements.get_bank_verification(account_id).data
-        assert "PDF" in pdf_response
+        assert "PDF" in str(pdf_response)
 
         params = GetStatementParams(s.id, "pdf")
         pdf_statement = client.statements.get(params).data
-        assert "PDF" in pdf_statement
-
-test_list_and_get_statements()
+        assert "PDF" in str(pdf_statement)
diff --git a/unit/api/statement_resource.py b/unit/api/statement_resource.py
@@ -8,29 +8,33 @@ class StatementResource(BaseResource):
     def __init__(self, configuration: Configuration):
         super().__init__("statements", configuration)
 
-    def get(self, params: GetStatementParams) -> Union[UnitResponse[str], UnitError]:
+    def get(self, params: GetStatementParams) -> Union[UnitResponse[bytes], UnitError]:
         parameters = {"language": params.language}
         if params.customer_id:
             parameters["filter[customerId]"] = params.customer_id
 
         response = super().get(f"{self.resource}/{params.statement_id}/{params.output_type}", parameters)
-        if response.status_code == 200:
-            return UnitResponse[str](response.text, None)
+
+        if super().is_20x(response.status_code):
+            return UnitResponse[bytes](response.content, None)
         else:
             return UnitError.from_json_api(response.json())
 
-    def get_bank_verification(self, account_id: str, include_proof_of_funds: Optional[bool] = False) -> Union[UnitResponse[str], UnitError]:
+    def get_bank_verification(self, account_id: str, include_proof_of_funds: Optional[bool] = False) -> \
+            Union[UnitResponse[bytes], UnitError]:
         response = super().get(f"{self.resource}/{account_id}/bank/pdf",
                                {"includeProofOfFunds": include_proof_of_funds})
-        if response.status_code == 200:
-            return UnitResponse[str](response.text, None)
+
+        if super().is_20x(response.status_code):
+            return UnitResponse[bytes](response.content, None)
         else:
             return UnitError.from_json_api(response.json())
 
     def list(self, params: ListStatementParams = None) -> Union[UnitResponse[List[StatementDTO]], UnitError]:
         params = params or ListStatementParams()
         response = super().get(self.resource, params.to_dict())
-        if response.status_code == 200:
+
+        if super().is_20x(response.status_code):
             data = response.json().get("data")
             return UnitResponse[StatementDTO](DtoDecoder.decode(data), None)
         else:
