From 4057e9c170267302d7b740980a91c7af26e044ef Mon Sep 17 00:00:00 2001
From: Bara Molova <b.molova@gmail.com>
Date: Mon, 5 Feb 2024 14:32:39 +0100
Subject: [PATCH] feat: add new retrieve_document_content method

---
 rossum_api/api_client.py                |  1 +
 rossum_api/elis_api_client.py           |  8 +++++
 rossum_api/elis_api_client_sync.py      |  7 +++++
 tests/elis_api_client/test_documents.py | 41 +++++++++++++++++++++++++
 4 files changed, 57 insertions(+)
 create mode 100644 tests/elis_api_client/test_documents.py

diff --git a/rossum_api/api_client.py b/rossum_api/api_client.py
index dce4f9a..812aea2 100644
--- a/rossum_api/api_client.py
+++ b/rossum_api/api_client.py
@@ -30,6 +30,7 @@ class Resource(Enum):
     Annotation = "annotations"
     Auth = "auth"
     Connector = "connectors"
+    Document = "documents"
     Group = "groups"
     Hook = "hooks"
     Inbox = "inboxes"
diff --git a/rossum_api/elis_api_client.py b/rossum_api/elis_api_client.py
index eeb1754..88c3c06 100644
--- a/rossum_api/elis_api_client.py
+++ b/rossum_api/elis_api_client.py
@@ -351,6 +351,14 @@ async def confirm_annotation(self, annotation_id: int) -> None:
             "POST", f"{Resource.Annotation.value}/{annotation_id}/confirm"
         )
 
+    # ##### DOCUMENTS #####
+    async def retrieve_document_content(self, document_id: int) -> bytes:
+        """https://elis.rossum.ai/api/docs/#document-content"""
+        document_content = await self._http_client.request(
+            "GET", url=f"{Resource.Document.value}/{document_id}/content"
+        )
+        return document_content.content
+
     # ##### WORKSPACES #####
     async def list_all_workspaces(
         self,
diff --git a/rossum_api/elis_api_client_sync.py b/rossum_api/elis_api_client_sync.py
index af73865..f191082 100644
--- a/rossum_api/elis_api_client_sync.py
+++ b/rossum_api/elis_api_client_sync.py
@@ -314,6 +314,13 @@ def confirm_annotation(self, annotation_id: int) -> None:
         """https://elis.rossum.ai/api/docs/#confirm-annotation"""
         self.event_loop.run_until_complete(self.elis_api_client.confirm_annotation(annotation_id))
 
+    # ##### DOCUMENTS #####
+    def retrieve_document_content(self, document_id: int) -> bytes:
+        """https://elis.rossum.ai/api/docs/#document-content"""
+        return self.event_loop.run_until_complete(
+            self.elis_api_client.retrieve_document_content(document_id)
+        )
+
     # ##### WORKSPACES #####
     def list_all_workspaces(
         self,
diff --git a/tests/elis_api_client/test_documents.py b/tests/elis_api_client/test_documents.py
new file mode 100644
index 0000000..8871cb4
--- /dev/null
+++ b/tests/elis_api_client/test_documents.py
@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import httpx
+import pytest
+
+from rossum_api.api_client import Resource
+
+
+@pytest.fixture
+def file_data() -> bytes:
+    with open("tests/data/sample_invoice.pdf", "rb") as f:
+        return f.read()
+
+
+@pytest.mark.asyncio
+class TestDocuments:
+    async def test_retrieve_document_content(self, elis_client, file_data):
+        client, http_client = elis_client
+        http_client.request.return_value = httpx.Response(status_code=200, content=file_data)
+
+        document_id = 123
+        result = await client.retrieve_document_content(document_id=document_id)
+
+        http_client.request.assert_called_with(
+            "GET", url=f"{Resource.Document.value}/{document_id}/content"
+        )
+        assert result == file_data
+
+
+class TestDocumentsSync:
+    def test_retrieve_document_content(self, elis_client_sync, file_data):
+        client, http_client = elis_client_sync
+        http_client.request.return_value = httpx.Response(status_code=200, content=file_data)
+
+        document_id = 123
+        result = client.retrieve_document_content(document_id=document_id)
+
+        http_client.request.assert_called_with(
+            "GET", url=f"{Resource.Document.value}/{document_id}/content"
+        )
+        assert result == file_data