fixed response type for drive item

Author: shemogumbe

src/msgraph_core/tasks/large_file_upload.py

@@ -4,13 +4,16 @@
 from asyncio import Future
 from datetime import datetime, timedelta, timezone
 import logging
+import httpx
 
 from kiota_abstractions.serialization.parsable import Parsable
 from kiota_abstractions.method import Method
 from kiota_abstractions.headers_collection import HeadersCollection
 from kiota_abstractions.request_information import RequestInformation
+from kiota_abstractions.native_response_handler import NativeResponseHandler
 from kiota_abstractions.serialization.additional_data_holder import AdditionalDataHolder
 from kiota_abstractions.serialization.parsable_factory import ParsableFactory
+from msgraph.generated.models.attachment_item import AttachmentItem
 
 from kiota_abstractions.request_adapter import RequestAdapter
 
@@ -25,7 +28,7 @@ def __init__(
         request_adapter: RequestAdapter,
         stream: BytesIO,
         parsable_factory: Optional[ParsableFactory] = None,
-        max_chunk_size: int = 5 * 1024 * 1024
+        max_chunk_size: int = 3 * 1024 * 1024  #5 - 409600
     ):
         self._upload_session = upload_session
         self._request_adapter = request_adapter
@@ -101,40 +104,45 @@ async def upload(self, after_chunk_upload: Optional[Callable] = None):
         end = min(int(range_parts[0]) + self.max_chunk_size - 1, self.file_size)
         uploaded_range = [range_parts[0], end]
         response = None
+        upload_result = UploadResult()
 
         while self.chunks >= 0:
             session = process_next
-            print(f"Chunks for upload  : {self.chunks}")
             if self.chunks == 0:
                 # last chunk
-                print(f"Last chunk: {self.chunks} upload stated")
                 response = await self.last_chunk(self.stream)
-                print("Last chunk response: received")
-
-            try:
-                lfu_session: LargeFileUploadSession = session  # type: ignore
-                if lfu_session is None:
-                    continue
-                next_range = lfu_session.next_expected_ranges
-                old_url = self.get_validated_upload_url(self.upload_session)
-                lfu_session.upload_url = old_url
-                if self.on_chunk_upload_complete is not None:
-                    self.on_chunk_upload_complete(uploaded_range)
-                if not next_range:
-                    continue
-                range_parts = str(next_range[0]).split("-")
-                end = min(int(range_parts[0]) + self.max_chunk_size, self.file_size)
-                uploaded_range = [range_parts[0], end]
-                self.next_range = next_range[0] + "-"
-                process_next = await self.next_chunk(self.stream)
-
-            except Exception as error:
-                logging.error("Error uploading chunk  %s", error)
-            finally:
-                self.chunks -= 1
+                # upload_result.location = response.headers[
+                #     'Location']  # to be fixed for attachement item Response
+
+            if isinstance(session, LargeFileUploadSession):
+                try:
+                    lfu_session: LargeFileUploadSession = session  # type: ignore
+                    if lfu_session is None:
+                        continue
+                    next_range = lfu_session.next_expected_ranges
+                    old_url = self.get_validated_upload_url(self.upload_session)
+                    lfu_session.upload_url = old_url
+                    if self.on_chunk_upload_complete is not None:
+                        self.on_chunk_upload_complete(uploaded_range)
+                    if not next_range:
+                        continue
+                    range_parts = str(next_range[0]).split("-")
+                    end = min(int(range_parts[0]) + self.max_chunk_size, self.file_size)
+                    uploaded_range = [range_parts[0], end]
+                    self.next_range = next_range[0] + "-"
+                    process_next = await self.next_chunk(self.stream)
+
+                except Exception as error:
+                    logging.error("Error uploading chunk  %s", error)
+                finally:
+                    self.chunks -= 1
+            else:
+                response = session
+                logging.info(f"Response headers: {response.headers}")
+                logging.info(f"Response content: {response.content}")
+                break
         upload_result = UploadResult()
         upload_result.item_response = response
-        upload_result.location = self.upload_session.upload_url
         return upload_result
 
     @property
@@ -177,6 +185,13 @@ async def next_chunk(self, file: BytesIO, range_start: int = 0, range_end: int =
         info.headers.try_add("Content-Type", "application/octet-stream")
         info.set_stream_content(bytes(chunk_data))
         error_map: Dict[str, int] = {}
+        if self.factory is AttachmentItem:
+            headers = {key: ', '.join(value) for key, value in info.headers.get_all().items()}
+            async with httpx.AsyncClient() as client:
+                response = await client.put(
+                    self.upload_session.upload_url, headers=headers, data=info.content
+                )
+                return response
         parsable_factory = LargeFileUploadSession
         return await self.request_adapter.send_async(info, parsable_factory, error_map)
 
@@ -218,6 +233,14 @@ async def last_chunk(
         info.headers.try_add("Content-Type", "application/octet-stream")
         info.set_stream_content(bytes(chunk_data))
         error_map: Dict[str, int] = {}
+        if self.factory is AttachmentItem:
+            headers = {key: ', '.join(value) for key, value in info.headers.get_all().items()}
+
+            async with httpx.AsyncClient() as client:
+                response = await client.put(
+                    self.upload_session.upload_url, headers=headers, data=info.content
+                )
+                return response
         parsable_factory = self.factory or parsable_factory
         return await self.request_adapter.send_async(info, parsable_factory, error_map)