diff --git a/google/cloud/documentai_toolbox/constants.py b/google/cloud/documentai_toolbox/constants.py index 5facde23..278d5446 100644 --- a/google/cloud/documentai_toolbox/constants.py +++ b/google/cloud/documentai_toolbox/constants.py @@ -21,8 +21,8 @@ FILE_CHECK_REGEX = r"(.*[.].*$)" -# https://cloud.google.com/document-ai/quotas -BATCH_MAX_FILES = 50 +# https://cloud.google.com/document-ai/quotas#content_limits +BATCH_MAX_FILES = 1000 # 1GB in Bytes BATCH_MAX_FILE_SIZE = 1073741824 BATCH_MAX_REQUESTS = 5 @@ -30,7 +30,8 @@ # https://cloud.google.com/document-ai/docs/file-types VALID_MIME_TYPES = { "application/pdf", - "image/bmp" "image/gif", + "image/bmp", + "image/gif", "image/jpeg", "image/png", "image/tiff", diff --git a/tests/unit/test_utilities.py b/tests/unit/test_utilities.py index fefff366..fd9335d8 100644 --- a/tests/unit/test_utilities.py +++ b/tests/unit/test_utilities.py @@ -350,10 +350,10 @@ def test_create_batches_with_3_documents(mock_storage, capfd): def test_create_batches_with_invalid_batch_size(): with pytest.raises( ValueError, - match="Batch size must be less than 50. You provided 51.", + match="Batch size must be less than 1000. You provided 1001.", ): utilities.create_batches( - gcs_bucket_name=test_bucket, gcs_prefix=test_prefix, batch_size=51 + gcs_bucket_name=test_bucket, gcs_prefix=test_prefix, batch_size=1001 ) @@ -373,7 +373,7 @@ def test_create_batches_with_large_folder(mock_storage, capfd): client.list_blobs.return_value = mock_blobs actual = utilities.create_batches( - gcs_bucket_name=test_bucket, gcs_prefix=test_prefix + gcs_bucket_name=test_bucket, gcs_prefix=test_prefix, batch_size=50 ) mock_storage.Client.assert_called_once()