|
16 | 16 | import time |
17 | 17 | from datetime import datetime, timezone |
18 | 18 | from io import BytesIO |
19 | | -from typing import Any |
| 19 | +from typing import Any, BinaryIO |
20 | 20 | from unittest.mock import Mock, patch |
21 | 21 |
|
22 | 22 | from httpx import URL, Request |
23 | 23 |
|
24 | | -from cognite.client.data_classes import Event, Row |
| 24 | +from cognite.client.data_classes import Event, FileMetadata, Row |
| 25 | +from cognite.client.data_classes.data_modeling.extractor_extensions.v1 import ( |
| 26 | + CogniteExtractorFileApply, |
| 27 | +) |
| 28 | +from cognite.client.testing import CogniteClientMock |
25 | 29 | from cognite.extractorutils.uploader import ( |
26 | 30 | EventUploadQueue, |
27 | 31 | IOFileUploadQueue, |
@@ -300,3 +304,40 @@ def test_mock_private_link_upload(MockCogniteClient: Mock) -> None: |
300 | 304 | response: Request = queue._get_file_upload_request(mock_download_url, mock_stream, len(bytes_)) |
301 | 305 |
|
302 | 306 | assert response.url.netloc == base_url.netloc |
| 307 | + |
| 308 | + |
| 309 | +@patch("cognite.client.CogniteClient") |
| 310 | +def test_manually_set_is_uploaded(mock_client: CogniteClientMock) -> None: |
| 311 | + mock_client.config.max_workers = 4 |
| 312 | + |
| 313 | + queue = IOFileUploadQueue( |
| 314 | + cdf_client=mock_client, |
| 315 | + max_queue_size=10, |
| 316 | + ) |
| 317 | + |
| 318 | + file_apply = CogniteExtractorFileApply( |
| 319 | + external_id="test-file", |
| 320 | + name="test-file.txt", |
| 321 | + space="test-space", |
| 322 | + ) |
| 323 | + |
| 324 | + # Record the initial state of is_uploaded |
| 325 | + initial_is_uploaded = file_apply.is_uploaded |
| 326 | + |
| 327 | + def read_file() -> BinaryIO: |
| 328 | + return BytesIO(b"test content") |
| 329 | + |
| 330 | + # Mock the upload methods |
| 331 | + with patch.object(queue, "_upload_bytes"), patch.object(queue, "_upload_only_metadata") as mock_upload_metadata: |
| 332 | + mock_file_metadata = Mock(spec=FileMetadata) |
| 333 | + mock_file_metadata.mime_type = "text/plain" |
| 334 | + mock_upload_metadata.return_value = (mock_file_metadata, "https://upload.url") |
| 335 | + |
| 336 | + with queue: |
| 337 | + queue.add_io_to_upload_queue(file_apply, read_file) |
| 338 | + queue.upload() |
| 339 | + |
| 340 | + # Verify that is_uploaded was NOT manually changed by the uploader |
| 341 | + assert file_apply.is_uploaded == initial_is_uploaded, ( |
| 342 | + "is_uploaded should not be manually changed by the uploader, it should be managed by the SDK" |
| 343 | + ) |
0 commit comments