From 70c6ab91df7b424b981d245fbacc2cdbe2905cdd Mon Sep 17 00:00:00 2001 From: VishalS-Microsoft Date: Tue, 31 Mar 2026 12:33:26 +0530 Subject: [PATCH 01/16] feat: Add weekly schedule for Azure Template validation,split azure-dev to azd-template-validation --- .github/workflows/azd-template-validation.yml | 46 +++++++++ .github/workflows/azure-dev.yml | 95 ++++++++++++------- 2 files changed, 105 insertions(+), 36 deletions(-) create mode 100644 .github/workflows/azd-template-validation.yml diff --git a/.github/workflows/azd-template-validation.yml b/.github/workflows/azd-template-validation.yml new file mode 100644 index 00000000..f59a8cbc --- /dev/null +++ b/.github/workflows/azd-template-validation.yml @@ -0,0 +1,46 @@ +name: AZD Template Validation + +on: + schedule: + - cron: '30 1 * * 4' # Every Thursday at 7:00 AM IST (1:30 AM UTC) + workflow_dispatch: + push: + branches: + - psl-weeklyschedule-codmod + +permissions: + contents: read + id-token: write + pull-requests: write + +jobs: + template_validation: + runs-on: ubuntu-latest + name: azd template validation + environment: production + steps: + - name: Checkout code + uses: actions/checkout@v6 + + - name: Validate Azure Template + id: validation + uses: microsoft/template-validation-action@v0.4.3 + with: + validateAzd: ${{ vars.TEMPLATE_VALIDATE_AZD }} + validateTests: ${{ vars.TEMPLATE_VALIDATE_TESTS }} + useDevContainer: ${{ vars.TEMPLATE_USE_DEV_CONTAINER }} + + env: + AZURE_CLIENT_ID: ${{ secrets.AZURE_CLIENT_ID }} + AZURE_TENANT_ID: ${{ secrets.AZURE_TENANT_ID }} + AZURE_SUBSCRIPTION_ID: ${{ secrets.AZURE_SUBSCRIPTION_ID }} + AZURE_ENV_NAME: ${{ secrets.AZURE_ENV_NAME }} + AZURE_LOCATION: ${{ secrets.AZURE_LOCATION }} + AZURE_ENV_AI_SERVICE_LOCATION: ${{ secrets.AZURE_AI_DEPLOYMENT_LOCATION || secrets.AZURE_LOCATION }} + AZURE_ENV_MODEL_CAPACITY: 1 + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + AZURE_DEV_COLLECT_TELEMETRY: ${{ vars.AZURE_DEV_COLLECT_TELEMETRY }} + + - name: Print result + shell: bash + run: cat "${{ steps.validation.outputs.resultFile }}" \ No newline at end of file diff --git a/.github/workflows/azure-dev.yml b/.github/workflows/azure-dev.yml index 7e5a6f57..c3a4793d 100644 --- a/.github/workflows/azure-dev.yml +++ b/.github/workflows/azure-dev.yml @@ -1,37 +1,60 @@ -name: Azure Template Validation -on: - workflow_dispatch: - -permissions: - contents: read - id-token: write - pull-requests: write -jobs: - template_validation_job: - runs-on: ubuntu-latest +name: Azure Dev Deploy + +on: + workflow_dispatch: + +permissions: + contents: read + id-token: write + +jobs: + deploy: + runs-on: ubuntu-latest environment: production - name: Template validation - steps: - # Step 1: Checkout the code from your repository - - name: Checkout code - uses: actions/checkout@v6 - # Step 2: Validate the Azure template using microsoft/template-validation-action - - name: Validate Azure Template - uses: microsoft/template-validation-action@v0.4.3 - with: - validateAzd: true - useDevContainer: false - id: validation - env: - AZURE_CLIENT_ID: ${{ secrets.AZURE_CLIENT_ID }} - AZURE_TENANT_ID: ${{ secrets.AZURE_TENANT_ID }} - AZURE_SUBSCRIPTION_ID: ${{ secrets.AZURE_SUBSCRIPTION_ID }} - AZURE_ENV_NAME: ${{ secrets.AZURE_ENV_NAME }} - AZURE_LOCATION: ${{ secrets.AZURE_LOCATION }} - AZURE_AI_DEPLOYMENT_LOCATION : ${{ secrets.AZURE_AI_DEPLOYMENT_LOCATION }} - AZURE_ENV_MODEL_CAPACITY : 1 - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - AZURE_DEV_COLLECT_TELEMETRY: ${{ vars.AZURE_DEV_COLLECT_TELEMETRY }} - # Step 3: Print the result of the validation - - name: Print result - run: cat ${{ steps.validation.outputs.resultFile }} + env: + AZURE_CLIENT_ID: ${{ secrets.AZURE_CLIENT_ID }} + AZURE_TENANT_ID: ${{ secrets.AZURE_TENANT_ID }} + AZURE_SUBSCRIPTION_ID: ${{ secrets.AZURE_SUBSCRIPTION_ID }} + AZURE_ENV_NAME: ${{ secrets.AZURE_ENV_NAME }} + AZURE_LOCATION: ${{ secrets.AZURE_LOCATION }} + AZURE_AI_DEPLOYMENT_LOCATION: ${{ secrets.AZURE_AI_DEPLOYMENT_LOCATION || secrets.AZURE_LOCATION }} + AZURE_ENV_MODEL_CAPACITY: 1 + AZURE_DEV_COLLECT_TELEMETRY: ${{ vars.AZURE_DEV_COLLECT_TELEMETRY }} + steps: + - name: Checkout code + uses: actions/checkout@v6 + + - name: Install azd + uses: Azure/setup-azd@v2 + + - name: Login to Azure + uses: azure/login@v2 + with: + client-id: ${{ secrets.AZURE_CLIENT_ID }} + tenant-id: ${{ secrets.AZURE_TENANT_ID }} + subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }} + + - name: Login to AZD + shell: bash + run: | + azd auth login \ + --client-id "$AZURE_CLIENT_ID" \ + --federated-credential-provider "github" \ + --tenant-id "$AZURE_TENANT_ID" + + - name: Provision and deploy + shell: bash + run: | + set -e + + if ! azd env select "$AZURE_ENV_NAME"; then + azd env new "$AZURE_ENV_NAME" --subscription "$AZURE_SUBSCRIPTION_ID" --location "$AZURE_LOCATION" --no-prompt + fi + + azd config set defaults.subscription "$AZURE_SUBSCRIPTION_ID" + azd env set AZURE_SUBSCRIPTION_ID "$AZURE_SUBSCRIPTION_ID" + azd env set AZURE_LOCATION "$AZURE_LOCATION" + azd env set AZURE_ENV_AI_SERVICE_LOCATION "${AZURE_AI_DEPLOYMENT_LOCATION:-$AZURE_LOCATION}" + azd env set AZURE_ENV_MODEL_CAPACITY "$AZURE_ENV_MODEL_CAPACITY" + + azd up --no-prompt From 39c5076cd7e767bfdef77ed12b9e1c4dbb7f39a9 Mon Sep 17 00:00:00 2001 From: VishalS-Microsoft Date: Tue, 31 Mar 2026 13:07:20 +0530 Subject: [PATCH 02/16] feat: Enable push trigger for Azure Dev Deploy workflow on psl-weeklyschedule-codmod branch --- .github/workflows/azure-dev.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/azure-dev.yml b/.github/workflows/azure-dev.yml index c3a4793d..e7db5cf0 100644 --- a/.github/workflows/azure-dev.yml +++ b/.github/workflows/azure-dev.yml @@ -2,6 +2,9 @@ name: Azure Dev Deploy on: workflow_dispatch: + push: + branches: + - psl-weeklyschedule-codmod permissions: contents: read From 013bd7266904aa6b24f67146d54b1d5245ae4ecd Mon Sep 17 00:00:00 2001 From: VishalS-Microsoft Date: Wed, 1 Apr 2026 11:37:24 +0530 Subject: [PATCH 03/16] feat: Remove push trigger for psl-weeklyschedule-codmod branch in workflow files --- .github/workflows/azd-template-validation.yml | 3 --- .github/workflows/azure-dev.yml | 3 --- 2 files changed, 6 deletions(-) diff --git a/.github/workflows/azd-template-validation.yml b/.github/workflows/azd-template-validation.yml index f59a8cbc..25eb8a1c 100644 --- a/.github/workflows/azd-template-validation.yml +++ b/.github/workflows/azd-template-validation.yml @@ -4,9 +4,6 @@ on: schedule: - cron: '30 1 * * 4' # Every Thursday at 7:00 AM IST (1:30 AM UTC) workflow_dispatch: - push: - branches: - - psl-weeklyschedule-codmod permissions: contents: read diff --git a/.github/workflows/azure-dev.yml b/.github/workflows/azure-dev.yml index e7db5cf0..c3a4793d 100644 --- a/.github/workflows/azure-dev.yml +++ b/.github/workflows/azure-dev.yml @@ -2,9 +2,6 @@ name: Azure Dev Deploy on: workflow_dispatch: - push: - branches: - - psl-weeklyschedule-codmod permissions: contents: read From db80ea72ac8ef487c2342be08b7a6b9040f63240 Mon Sep 17 00:00:00 2001 From: VishalS-Microsoft Date: Thu, 2 Apr 2026 19:12:13 +0530 Subject: [PATCH 04/16] feat: Add push trigger and timestamp setting for Azure workflows on psl-weeklyschedule-codmod branch --- .github/workflows/azd-template-validation.yml | 8 +++++++- .github/workflows/azure-dev.yml | 8 ++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/.github/workflows/azd-template-validation.yml b/.github/workflows/azd-template-validation.yml index 25eb8a1c..375dd539 100644 --- a/.github/workflows/azd-template-validation.yml +++ b/.github/workflows/azd-template-validation.yml @@ -4,6 +4,9 @@ on: schedule: - cron: '30 1 * * 4' # Every Thursday at 7:00 AM IST (1:30 AM UTC) workflow_dispatch: + push: + branches: + - psl-weeklyschedule-codmod permissions: contents: read @@ -19,6 +22,9 @@ jobs: - name: Checkout code uses: actions/checkout@v6 + - name: Set timestamp + run: echo "HHMM=$(date -u +'%H%M')" >> $GITHUB_ENV + - name: Validate Azure Template id: validation uses: microsoft/template-validation-action@v0.4.3 @@ -31,7 +37,7 @@ jobs: AZURE_CLIENT_ID: ${{ secrets.AZURE_CLIENT_ID }} AZURE_TENANT_ID: ${{ secrets.AZURE_TENANT_ID }} AZURE_SUBSCRIPTION_ID: ${{ secrets.AZURE_SUBSCRIPTION_ID }} - AZURE_ENV_NAME: ${{ secrets.AZURE_ENV_NAME }} + AZURE_ENV_NAME: azd-${{ secrets.AZURE_ENV_NAME }}-${{ env.HHMM }} AZURE_LOCATION: ${{ secrets.AZURE_LOCATION }} AZURE_ENV_AI_SERVICE_LOCATION: ${{ secrets.AZURE_AI_DEPLOYMENT_LOCATION || secrets.AZURE_LOCATION }} AZURE_ENV_MODEL_CAPACITY: 1 diff --git a/.github/workflows/azure-dev.yml b/.github/workflows/azure-dev.yml index c3a4793d..64aa3d5e 100644 --- a/.github/workflows/azure-dev.yml +++ b/.github/workflows/azure-dev.yml @@ -2,6 +2,9 @@ name: Azure Dev Deploy on: workflow_dispatch: + push: + branches: + - psl-weeklyschedule-codmod permissions: contents: read @@ -24,6 +27,11 @@ jobs: - name: Checkout code uses: actions/checkout@v6 + - name: Set timestamp and env name + run: | + HHMM=$(date -u +'%H%M') + echo "AZURE_ENV_NAME=azd-${{ vars.AZURE_ENV_NAME }}-${HHMM}" >> $GITHUB_ENV + - name: Install azd uses: Azure/setup-azd@v2 From 5e5b2daa2c4754f40a63e24ab0abf965584d928b Mon Sep 17 00:00:00 2001 From: VishalS-Microsoft Date: Thu, 2 Apr 2026 19:53:00 +0530 Subject: [PATCH 05/16] feat: Remove push trigger for psl-weeklyschedule-codmod branch in workflow files --- .github/workflows/azd-template-validation.yml | 3 --- .github/workflows/azure-dev.yml | 3 --- 2 files changed, 6 deletions(-) diff --git a/.github/workflows/azd-template-validation.yml b/.github/workflows/azd-template-validation.yml index 375dd539..397607fe 100644 --- a/.github/workflows/azd-template-validation.yml +++ b/.github/workflows/azd-template-validation.yml @@ -4,9 +4,6 @@ on: schedule: - cron: '30 1 * * 4' # Every Thursday at 7:00 AM IST (1:30 AM UTC) workflow_dispatch: - push: - branches: - - psl-weeklyschedule-codmod permissions: contents: read diff --git a/.github/workflows/azure-dev.yml b/.github/workflows/azure-dev.yml index 64aa3d5e..2b887434 100644 --- a/.github/workflows/azure-dev.yml +++ b/.github/workflows/azure-dev.yml @@ -2,9 +2,6 @@ name: Azure Dev Deploy on: workflow_dispatch: - push: - branches: - - psl-weeklyschedule-codmod permissions: contents: read From d7c67796403549a3ec0424d60d500158d107bdc4 Mon Sep 17 00:00:00 2001 From: Pavan-Microsoft Date: Fri, 3 Apr 2026 17:30:51 +0530 Subject: [PATCH 06/16] Enhance CosmosDBClient and DatabaseFactory for improved error handling and concurrency - Added asyncio support and a lock mechanism in DatabaseFactory to ensure thread safety. - Implemented retry logic with backoff for reading existing batches in CosmosDBClient. - Updated batch service to handle existing batch entries more efficiently. --- src/backend/common/database/cosmosdb.py | 37 ++++++++++++++----- .../common/database/database_factory.py | 31 ++++++++++------ src/backend/common/services/batch_service.py | 8 ++-- 3 files changed, 52 insertions(+), 24 deletions(-) diff --git a/src/backend/common/database/cosmosdb.py b/src/backend/common/database/cosmosdb.py index 7ff043f4..033de53a 100644 --- a/src/backend/common/database/cosmosdb.py +++ b/src/backend/common/database/cosmosdb.py @@ -1,3 +1,4 @@ +import asyncio from datetime import datetime, timezone from typing import Dict, List, Optional from uuid import UUID, uuid4 @@ -5,7 +6,8 @@ from azure.cosmos.aio import CosmosClient from azure.cosmos.aio._database import DatabaseProxy from azure.cosmos.exceptions import ( - CosmosResourceExistsError + CosmosResourceExistsError, + CosmosResourceNotFoundError, ) from common.database.database_base import DatabaseBase @@ -85,9 +87,21 @@ async def create_batch(self, user_id: str, batch_id: UUID) -> BatchRecord: await self.batch_container.create_item(body=batch.dict()) return batch except CosmosResourceExistsError: - self.logger.info(f"Batch with ID {batch_id} already exists") - batchexists = await self.get_batch(user_id, str(batch_id)) - return batchexists + self.logger.info("Batch already exists, reading existing record", batch_id=str(batch_id)) + # Retry read with backoff to handle replication lag after 409 conflict + for attempt in range(3): + try: + batchexists = await self.batch_container.read_item( + item=str(batch_id), partition_key=str(batch_id) + ) + if batchexists.get("user_id") != user_id: + raise ValueError("Batch belongs to a different user") + return batchexists + except CosmosResourceNotFoundError: + if attempt < 2: + await asyncio.sleep(0.5 * (attempt + 1)) + else: + return batch except Exception as e: self.logger.error("Failed to create batch", error=str(e)) @@ -158,7 +172,7 @@ async def get_batch(self, user_id: str, batch_id: str) -> Optional[Dict]: ] batch = None async for item in self.batch_container.query_items( - query=query, parameters=params + query=query, parameters=params, partition_key=batch_id ): batch = item @@ -173,7 +187,7 @@ async def get_file(self, file_id: str) -> Optional[Dict]: params = [{"name": "@file_id", "value": file_id}] file_entry = None async for item in self.file_container.query_items( - query=query, parameters=params + query=query, parameters=params, partition_key=file_id ): file_entry = item return file_entry @@ -209,7 +223,7 @@ async def get_batch_from_id(self, batch_id: str) -> Dict: batch = None # Store the batch async for item in self.batch_container.query_items( - query=query, parameters=params + query=query, parameters=params, partition_key=batch_id ): batch = item # Assign the batch to the variable @@ -335,11 +349,14 @@ async def add_file_log( raise async def update_batch_entry( - self, batch_id: str, user_id: str, status: ProcessStatus, file_count: int + self, batch_id: str, user_id: str, status: ProcessStatus, file_count: int, + existing_batch: Optional[Dict] = None ): - """Update batch status.""" + """Update batch status. If existing_batch is provided, skip the re-fetch.""" try: - batch = await self.get_batch(user_id, batch_id) + batch = existing_batch + if batch is None: + batch = await self.get_batch(user_id, batch_id) if not batch: raise ValueError("Batch not found") diff --git a/src/backend/common/database/database_factory.py b/src/backend/common/database/database_factory.py index c2f7de9d..4a6c84fc 100644 --- a/src/backend/common/database/database_factory.py +++ b/src/backend/common/database/database_factory.py @@ -9,25 +9,34 @@ class DatabaseFactory: _instance: Optional[DatabaseBase] = None + _lock: asyncio.Lock = asyncio.Lock() _logger = AppLogger("DatabaseFactory") @staticmethod async def get_database(): + if DatabaseFactory._instance is not None: + return DatabaseFactory._instance - config = Config() # Create an instance of Config + async with DatabaseFactory._lock: + # Double-check after acquiring the lock + if DatabaseFactory._instance is not None: + return DatabaseFactory._instance - cosmos_db_client = CosmosDBClient( - endpoint=config.cosmosdb_endpoint, - credential=config.get_azure_credentials(), - database_name=config.cosmosdb_database, - batch_container=config.cosmosdb_batch_container, - file_container=config.cosmosdb_file_container, - log_container=config.cosmosdb_log_container, - ) + config = Config() # Create an instance of Config - await cosmos_db_client.initialize_cosmos() + cosmos_db_client = CosmosDBClient( + endpoint=config.cosmosdb_endpoint, + credential=config.get_azure_credentials(), + database_name=config.cosmosdb_database, + batch_container=config.cosmosdb_batch_container, + file_container=config.cosmosdb_file_container, + log_container=config.cosmosdb_log_container, + ) - return cosmos_db_client + await cosmos_db_client.initialize_cosmos() + + DatabaseFactory._instance = cosmos_db_client + return cosmos_db_client # Local testing of config and code diff --git a/src/backend/common/services/batch_service.py b/src/backend/common/services/batch_service.py index f467820a..8e54eeda 100644 --- a/src/backend/common/services/batch_service.py +++ b/src/backend/common/services/batch_service.py @@ -287,8 +287,8 @@ async def upload_file_to_batch(self, batch_id: str, user_id: str, file: UploadFi ) # Create file entry - await self.database.add_file(batch_id, file_id, file.filename, blob_path) - file_record = await self.database.get_file(file_id) + file_record_obj = await self.database.add_file(batch_id, file_id, file.filename, blob_path) + file_record = file_record_obj.dict() if hasattr(file_record_obj, 'dict') else file_record_obj await self.database.add_file_log( UUID(file_id), @@ -307,6 +307,7 @@ async def upload_file_to_batch(self, batch_id: str, user_id: str, file: UploadFi user_id, ProcessStatus.READY_TO_PROCESS, batch["file_count"], + existing_batch=batch, ) # Return response return {"batch": batch, "file": file_record} @@ -317,7 +318,8 @@ async def upload_file_to_batch(self, batch_id: str, user_id: str, file: UploadFi batch.file_count = len(files) batch.updated_at = datetime.utcnow().isoformat() await self.database.update_batch_entry( - batch_id, user_id, ProcessStatus.READY_TO_PROCESS, batch.file_count + batch_id, user_id, ProcessStatus.READY_TO_PROCESS, batch.file_count, + existing_batch=batch.dict(), ) # Return response return {"batch": batch, "file": file_record} From a6b843a828fac1629c11456e543c592897265d40 Mon Sep 17 00:00:00 2001 From: Pavan-Microsoft Date: Fri, 3 Apr 2026 19:03:06 +0530 Subject: [PATCH 07/16] Refactor test_create_batch_exists to mock read_item instead of get_batch for existing batch records --- .../backend/common/database/cosmosdb_test.py | 33 ++++++++++--------- 1 file changed, 18 insertions(+), 15 deletions(-) diff --git a/src/tests/backend/common/database/cosmosdb_test.py b/src/tests/backend/common/database/cosmosdb_test.py index 3405d85e..46cfe206 100644 --- a/src/tests/backend/common/database/cosmosdb_test.py +++ b/src/tests/backend/common/database/cosmosdb_test.py @@ -158,31 +158,34 @@ async def test_create_batch_exists(cosmos_db_client, mocker): user_id = "user_1" batch_id = uuid4() - # Mock container creation and get_batch + # Mock container creation and read_item mock_batch_container = mock.MagicMock() mocker.patch.object(cosmos_db_client, 'batch_container', mock_batch_container) mock_batch_container.create_item = AsyncMock(side_effect=CosmosResourceExistsError) - # Mock the get_batch method - mock_get_batch = AsyncMock(return_value=BatchRecord( - batch_id=batch_id, - user_id=user_id, - file_count=0, - created_at=datetime.now(timezone.utc), - updated_at=datetime.now(timezone.utc), - status=ProcessStatus.READY_TO_PROCESS - )) - mocker.patch.object(cosmos_db_client, 'get_batch', mock_get_batch) + # Mock read_item to return the existing batch record + existing_batch = { + "id": str(batch_id), + "batch_id": str(batch_id), + "user_id": user_id, + "file_count": 0, + "created_at": datetime.now(timezone.utc).isoformat(), + "updated_at": datetime.now(timezone.utc).isoformat(), + "status": ProcessStatus.READY_TO_PROCESS, + } + mock_batch_container.read_item = AsyncMock(return_value=existing_batch) # Call the method batch = await cosmos_db_client.create_batch(user_id, batch_id) # Assert that batch was fetched (not created) due to already existing - assert batch.batch_id == batch_id - assert batch.user_id == user_id - assert batch.status == ProcessStatus.READY_TO_PROCESS + assert batch["batch_id"] == str(batch_id) + assert batch["user_id"] == user_id + assert batch["status"] == ProcessStatus.READY_TO_PROCESS - mock_get_batch.assert_called_once_with(user_id, str(batch_id)) + mock_batch_container.read_item.assert_called_once_with( + item=str(batch_id), partition_key=str(batch_id) + ) @pytest.mark.asyncio From 7815d5c5bedab67735a4e15426c0f10720155f05 Mon Sep 17 00:00:00 2001 From: Pavan-Microsoft Date: Fri, 3 Apr 2026 19:11:14 +0530 Subject: [PATCH 08/16] Enhance CosmosDBClient error logging and improve test assertions for existing batch records --- src/backend/common/database/cosmosdb.py | 9 +++++++-- src/tests/backend/common/database/cosmosdb_test.py | 6 +++--- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/backend/common/database/cosmosdb.py b/src/backend/common/database/cosmosdb.py index 033de53a..b7dec35c 100644 --- a/src/backend/common/database/cosmosdb.py +++ b/src/backend/common/database/cosmosdb.py @@ -95,13 +95,18 @@ async def create_batch(self, user_id: str, batch_id: UUID) -> BatchRecord: item=str(batch_id), partition_key=str(batch_id) ) if batchexists.get("user_id") != user_id: + self.logger.error("Batch belongs to a different user", batch_id=str(batch_id), existing_user=batchexists.get("user_id"), requesting_user=user_id) raise ValueError("Batch belongs to a different user") - return batchexists + self.logger.info("Returning existing batch record", batch_id=str(batch_id)) + return BatchRecord.fromdb(batchexists) except CosmosResourceNotFoundError: if attempt < 2: + self.logger.info("Batch read returned 404 after conflict, retrying", batch_id=str(batch_id), attempt=attempt + 1) await asyncio.sleep(0.5 * (attempt + 1)) else: - return batch + raise RuntimeError( + f"Batch {batch_id} already exists but could not be read after retries" + ) except Exception as e: self.logger.error("Failed to create batch", error=str(e)) diff --git a/src/tests/backend/common/database/cosmosdb_test.py b/src/tests/backend/common/database/cosmosdb_test.py index 46cfe206..a9483406 100644 --- a/src/tests/backend/common/database/cosmosdb_test.py +++ b/src/tests/backend/common/database/cosmosdb_test.py @@ -179,9 +179,9 @@ async def test_create_batch_exists(cosmos_db_client, mocker): batch = await cosmos_db_client.create_batch(user_id, batch_id) # Assert that batch was fetched (not created) due to already existing - assert batch["batch_id"] == str(batch_id) - assert batch["user_id"] == user_id - assert batch["status"] == ProcessStatus.READY_TO_PROCESS + assert batch.batch_id == batch_id + assert batch.user_id == user_id + assert batch.status == ProcessStatus.READY_TO_PROCESS mock_batch_container.read_item.assert_called_once_with( item=str(batch_id), partition_key=str(batch_id) From e481341509202a62fc1a1aec0a2edcf308cc14ae Mon Sep 17 00:00:00 2001 From: Pavan-Microsoft Date: Fri, 3 Apr 2026 19:49:33 +0530 Subject: [PATCH 09/16] Refactor DatabaseFactory to lazily initialize the lock and improve file record handling in BatchService --- src/backend/common/database/database_factory.py | 10 ++++++++-- src/backend/common/services/batch_service.py | 3 ++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/src/backend/common/database/database_factory.py b/src/backend/common/database/database_factory.py index 4a6c84fc..bbc84941 100644 --- a/src/backend/common/database/database_factory.py +++ b/src/backend/common/database/database_factory.py @@ -9,15 +9,21 @@ class DatabaseFactory: _instance: Optional[DatabaseBase] = None - _lock: asyncio.Lock = asyncio.Lock() + _lock: Optional[asyncio.Lock] = None _logger = AppLogger("DatabaseFactory") + @staticmethod + def _get_lock() -> asyncio.Lock: + if DatabaseFactory._lock is None: + DatabaseFactory._lock = asyncio.Lock() + return DatabaseFactory._lock + @staticmethod async def get_database(): if DatabaseFactory._instance is not None: return DatabaseFactory._instance - async with DatabaseFactory._lock: + async with DatabaseFactory._get_lock(): # Double-check after acquiring the lock if DatabaseFactory._instance is not None: return DatabaseFactory._instance diff --git a/src/backend/common/services/batch_service.py b/src/backend/common/services/batch_service.py index 8e54eeda..bbcb75e5 100644 --- a/src/backend/common/services/batch_service.py +++ b/src/backend/common/services/batch_service.py @@ -288,7 +288,8 @@ async def upload_file_to_batch(self, batch_id: str, user_id: str, file: UploadFi # Create file entry file_record_obj = await self.database.add_file(batch_id, file_id, file.filename, blob_path) - file_record = file_record_obj.dict() if hasattr(file_record_obj, 'dict') else file_record_obj + file_record_dict = getattr(file_record_obj, "dict", None) + file_record = file_record_dict() if callable(file_record_dict) else file_record_obj await self.database.add_file_log( UUID(file_id), From 6bdd50e032c8f8d2ed7baf0a8093d31424528458 Mon Sep 17 00:00:00 2001 From: Pavan-Microsoft Date: Fri, 3 Apr 2026 19:59:49 +0530 Subject: [PATCH 10/16] fix: avoid leaking batch existence on user mismatch (Copilot review) --- src/backend/common/database/cosmosdb.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/backend/common/database/cosmosdb.py b/src/backend/common/database/cosmosdb.py index b7dec35c..46c020d0 100644 --- a/src/backend/common/database/cosmosdb.py +++ b/src/backend/common/database/cosmosdb.py @@ -95,8 +95,8 @@ async def create_batch(self, user_id: str, batch_id: UUID) -> BatchRecord: item=str(batch_id), partition_key=str(batch_id) ) if batchexists.get("user_id") != user_id: - self.logger.error("Batch belongs to a different user", batch_id=str(batch_id), existing_user=batchexists.get("user_id"), requesting_user=user_id) - raise ValueError("Batch belongs to a different user") + self.logger.error("Batch belongs to a different user", batch_id=str(batch_id)) + raise CosmosResourceNotFoundError(message="Batch not found") self.logger.info("Returning existing batch record", batch_id=str(batch_id)) return BatchRecord.fromdb(batchexists) except CosmosResourceNotFoundError: From bda0369d63ee708695b402d9ddf4e6737cbf8188 Mon Sep 17 00:00:00 2001 From: Pavan-Microsoft Date: Fri, 3 Apr 2026 20:05:26 +0530 Subject: [PATCH 11/16] fix: update mock_query_items to accept additional kwargs for improved flexibility --- src/tests/backend/common/database/cosmosdb_test.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/tests/backend/common/database/cosmosdb_test.py b/src/tests/backend/common/database/cosmosdb_test.py index a9483406..d08fb015 100644 --- a/src/tests/backend/common/database/cosmosdb_test.py +++ b/src/tests/backend/common/database/cosmosdb_test.py @@ -407,7 +407,7 @@ async def test_get_batch(cosmos_db_client, mocker): } # We define the async generator function that will yield the expected batch - async def mock_query_items(query, parameters): + async def mock_query_items(query, parameters, **kwargs): yield expected_batch # Assign the async generator to query_items mock @@ -425,6 +425,7 @@ async def mock_query_items(query, parameters): {"name": "@batch_id", "value": batch_id}, {"name": "@user_id", "value": user_id}, ], + partition_key=batch_id, ) @@ -471,8 +472,8 @@ async def test_get_file(cosmos_db_client, mocker): "blob_path": "/path/to/file" } - # We define the async generator function that will yield the expected batch - async def mock_query_items(query, parameters): + # We define the async generator function that will yield the expected file + async def mock_query_items(query, parameters, **kwargs): yield expected_file # Assign the async generator to query_items mock @@ -597,7 +598,7 @@ async def test_get_batch_from_id(cosmos_db_client, mocker): } # Define the async generator function that will yield the expected batch - async def mock_query_items(query, parameters): + async def mock_query_items(query, parameters, **kwargs): yield expected_batch # Assign the async generator to query_items mock From 79d5964c27994df74ea3c7d08415af1f1aea8f5d Mon Sep 17 00:00:00 2001 From: Pavan-Microsoft Date: Mon, 6 Apr 2026 17:20:58 +0530 Subject: [PATCH 12/16] fix: update fallback behavior to check batch status instead of forcing navigation after 2 minutes --- src/frontend/src/pages/modernizationPage.tsx | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/frontend/src/pages/modernizationPage.tsx b/src/frontend/src/pages/modernizationPage.tsx index 53f128b5..3282dfcf 100644 --- a/src/frontend/src/pages/modernizationPage.tsx +++ b/src/frontend/src/pages/modernizationPage.tsx @@ -1052,16 +1052,16 @@ useEffect(() => { updateSummaryStatus(); } - // Ultimate fallback: If on page for 2+ minutes with no completion, force navigation + // Ultimate fallback: If on page for 2+ minutes with no completion, check batch status instead of force-navigating const timeSincePageLoad = Date.now() - pageLoadTime; if (timeSincePageLoad > 120000 && !allFilesCompleted && nonSummaryFiles.length > 0) { - console.log("Page loaded for 2+ minutes without completion, forcing navigation to batch view"); - navigate(`/batch-view/${batchId}`); + console.log("Page loaded for 2+ minutes without completion, checking batch status"); + updateSummaryStatus(); } }, 5000); // Check every 5 seconds return () => clearInterval(checkInactivity); - }, [lastActivityTime, files, allFilesCompleted, updateSummaryStatus, pageLoadTime, navigate, batchId]); + }, [lastActivityTime, files, allFilesCompleted, updateSummaryStatus, pageLoadTime, batchId]); useEffect(() => { From 27e2f53174b03bfa5aeff8e61c44274c68f35fdf Mon Sep 17 00:00:00 2001 From: Pavan-Microsoft Date: Mon, 6 Apr 2026 18:00:46 +0530 Subject: [PATCH 13/16] Revert "fix: update fallback behavior to check batch status instead of forcing navigation after 2 minutes" This reverts commit 79d5964c27994df74ea3c7d08415af1f1aea8f5d. --- src/frontend/src/pages/modernizationPage.tsx | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/frontend/src/pages/modernizationPage.tsx b/src/frontend/src/pages/modernizationPage.tsx index 3282dfcf..53f128b5 100644 --- a/src/frontend/src/pages/modernizationPage.tsx +++ b/src/frontend/src/pages/modernizationPage.tsx @@ -1052,16 +1052,16 @@ useEffect(() => { updateSummaryStatus(); } - // Ultimate fallback: If on page for 2+ minutes with no completion, check batch status instead of force-navigating + // Ultimate fallback: If on page for 2+ minutes with no completion, force navigation const timeSincePageLoad = Date.now() - pageLoadTime; if (timeSincePageLoad > 120000 && !allFilesCompleted && nonSummaryFiles.length > 0) { - console.log("Page loaded for 2+ minutes without completion, checking batch status"); - updateSummaryStatus(); + console.log("Page loaded for 2+ minutes without completion, forcing navigation to batch view"); + navigate(`/batch-view/${batchId}`); } }, 5000); // Check every 5 seconds return () => clearInterval(checkInactivity); - }, [lastActivityTime, files, allFilesCompleted, updateSummaryStatus, pageLoadTime, batchId]); + }, [lastActivityTime, files, allFilesCompleted, updateSummaryStatus, pageLoadTime, navigate, batchId]); useEffect(() => { From 5e765a8f47f35ddac78208697e1a4c711b2267e1 Mon Sep 17 00:00:00 2001 From: Pavan-Microsoft Date: Tue, 7 Apr 2026 13:18:34 +0530 Subject: [PATCH 14/16] fix: add dependency on aiProject for cognitive service deployments --- infra/modules/ai-foundry/dependencies.bicep | 3 +++ 1 file changed, 3 insertions(+) diff --git a/infra/modules/ai-foundry/dependencies.bicep b/infra/modules/ai-foundry/dependencies.bicep index f0d119cb..0d28aa1a 100644 --- a/infra/modules/ai-foundry/dependencies.bicep +++ b/infra/modules/ai-foundry/dependencies.bicep @@ -208,6 +208,9 @@ resource cognitiveService_deployments 'Microsoft.CognitiveServices/accounts/depl size: sku.?size family: sku.?family } + dependsOn: [ + aiProject + ] } ] From 99f379d48962f69d81482dfec46582d2252b666f Mon Sep 17 00:00:00 2001 From: Pavan-Microsoft Date: Tue, 7 Apr 2026 13:52:22 +0530 Subject: [PATCH 15/16] fix: update bicep version and template hashes in main.json --- infra/main.json | 64 +++++++++++++++++++++++++++---------------------- 1 file changed, 35 insertions(+), 29 deletions(-) diff --git a/infra/main.json b/infra/main.json index f98bb0c5..09e37656 100644 --- a/infra/main.json +++ b/infra/main.json @@ -5,8 +5,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "3093757051086668797" + "version": "0.42.1.51946", + "templateHash": "7222423000870488333" }, "name": "Modernize Your Code Solution Accelerator", "description": "CSA CTO Gold Standard Solution Accelerator for Modernize Your Code. \r\n" @@ -5052,8 +5052,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "8663094775498995429" + "version": "0.42.1.51946", + "templateHash": "3406526791248457038" } }, "definitions": { @@ -12895,11 +12895,11 @@ }, "dependsOn": [ "applicationInsights", - "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').storageBlob)]", "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').ods)]", - "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').agentSvc)]", "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').monitor)]", "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').oms)]", + "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').storageBlob)]", + "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').agentSvc)]", "dataCollectionEndpoint", "logAnalyticsWorkspace", "virtualNetwork" @@ -25611,8 +25611,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "17285204072656433491" + "version": "0.42.1.51946", + "templateHash": "16969185198334420434" }, "name": "AI Services and Project Module", "description": "This module creates an AI Services resource and an AI Foundry project within it. It supports private networking, OpenAI deployments, and role assignments." @@ -26952,8 +26952,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "5121330425393020264" + "version": "0.42.1.51946", + "templateHash": "4140498216793917924" } }, "definitions": { @@ -27910,7 +27910,10 @@ "raiPolicyName": "[tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'raiPolicyName')]", "versionUpgradeOption": "[tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'versionUpgradeOption')]" }, - "sku": "[coalesce(tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'sku'), createObject('name', parameters('sku'), 'capacity', tryGet(parameters('sku'), 'capacity'), 'tier', tryGet(parameters('sku'), 'tier'), 'size', tryGet(parameters('sku'), 'size'), 'family', tryGet(parameters('sku'), 'family')))]" + "sku": "[coalesce(tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'sku'), createObject('name', parameters('sku'), 'capacity', tryGet(parameters('sku'), 'capacity'), 'tier', tryGet(parameters('sku'), 'tier'), 'size', tryGet(parameters('sku'), 'size'), 'family', tryGet(parameters('sku'), 'family')))]", + "dependsOn": [ + "aiProject" + ] }, "cognitiveService_lock": { "condition": "[and(not(empty(coalesce(parameters('lock'), createObject()))), not(equals(tryGet(parameters('lock'), 'kind'), 'None')))]", @@ -28664,8 +28667,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "10989408486030617267" + "version": "0.42.1.51946", + "templateHash": "2422737205646151487" } }, "definitions": { @@ -28818,8 +28821,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "7933643033523871028" + "version": "0.42.1.51946", + "templateHash": "11911242767938607365" } }, "definitions": { @@ -29036,8 +29039,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "5121330425393020264" + "version": "0.42.1.51946", + "templateHash": "4140498216793917924" } }, "definitions": { @@ -29994,7 +29997,10 @@ "raiPolicyName": "[tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'raiPolicyName')]", "versionUpgradeOption": "[tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'versionUpgradeOption')]" }, - "sku": "[coalesce(tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'sku'), createObject('name', parameters('sku'), 'capacity', tryGet(parameters('sku'), 'capacity'), 'tier', tryGet(parameters('sku'), 'tier'), 'size', tryGet(parameters('sku'), 'size'), 'family', tryGet(parameters('sku'), 'family')))]" + "sku": "[coalesce(tryGet(coalesce(parameters('deployments'), createArray())[copyIndex()], 'sku'), createObject('name', parameters('sku'), 'capacity', tryGet(parameters('sku'), 'capacity'), 'tier', tryGet(parameters('sku'), 'tier'), 'size', tryGet(parameters('sku'), 'size'), 'family', tryGet(parameters('sku'), 'family')))]", + "dependsOn": [ + "aiProject" + ] }, "cognitiveService_lock": { "condition": "[and(not(empty(coalesce(parameters('lock'), createObject()))), not(equals(tryGet(parameters('lock'), 'kind'), 'None')))]", @@ -30748,8 +30754,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "10989408486030617267" + "version": "0.42.1.51946", + "templateHash": "2422737205646151487" } }, "definitions": { @@ -30902,8 +30908,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "7933643033523871028" + "version": "0.42.1.51946", + "templateHash": "11911242767938607365" } }, "definitions": { @@ -31917,8 +31923,8 @@ "dependsOn": [ "aiServices", "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').openAI)]", - "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').cognitiveServices)]", "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').aiServices)]", + "[format('avmPrivateDnsZones[{0}]', variables('dnsZoneIndex').cognitiveServices)]", "virtualNetwork" ] }, @@ -31974,8 +31980,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "3881415837167031634" + "version": "0.42.1.51946", + "templateHash": "522477461329004641" } }, "definitions": { @@ -40219,8 +40225,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "15962472891869337617" + "version": "0.42.1.51946", + "templateHash": "15355322017409205910" } }, "definitions": { @@ -44082,8 +44088,8 @@ "metadata": { "_generator": { "name": "bicep", - "version": "0.41.2.15936", - "templateHash": "17636459140972536078" + "version": "0.42.1.51946", + "templateHash": "4242598725709304634" } }, "definitions": { From fc8de18803da76806d067209eca2e6bd5585d44a Mon Sep 17 00:00:00 2001 From: Pavan-Microsoft Date: Tue, 7 Apr 2026 20:38:53 +0530 Subject: [PATCH 16/16] fix: enhance batch handling in CosmosDBClient and add tests for conflict scenarios --- src/backend/common/database/cosmosdb.py | 20 +++-- src/backend/common/services/batch_service.py | 2 +- .../backend/common/database/cosmosdb_test.py | 84 ++++++++++++++++++- 3 files changed, 95 insertions(+), 11 deletions(-) diff --git a/src/backend/common/database/cosmosdb.py b/src/backend/common/database/cosmosdb.py index 46c020d0..8c56286a 100644 --- a/src/backend/common/database/cosmosdb.py +++ b/src/backend/common/database/cosmosdb.py @@ -94,19 +94,21 @@ async def create_batch(self, user_id: str, batch_id: UUID) -> BatchRecord: batchexists = await self.batch_container.read_item( item=str(batch_id), partition_key=str(batch_id) ) - if batchexists.get("user_id") != user_id: - self.logger.error("Batch belongs to a different user", batch_id=str(batch_id)) - raise CosmosResourceNotFoundError(message="Batch not found") - self.logger.info("Returning existing batch record", batch_id=str(batch_id)) - return BatchRecord.fromdb(batchexists) except CosmosResourceNotFoundError: if attempt < 2: self.logger.info("Batch read returned 404 after conflict, retrying", batch_id=str(batch_id), attempt=attempt + 1) await asyncio.sleep(0.5 * (attempt + 1)) - else: - raise RuntimeError( - f"Batch {batch_id} already exists but could not be read after retries" - ) + continue + raise RuntimeError( + f"Batch {batch_id} already exists but could not be read after retries" + ) + + if batchexists.get("user_id") != user_id: + self.logger.error("Batch belongs to a different user", batch_id=str(batch_id)) + raise PermissionError("Batch not found") + + self.logger.info("Returning existing batch record", batch_id=str(batch_id)) + return BatchRecord.fromdb(batchexists) except Exception as e: self.logger.error("Failed to create batch", error=str(e)) diff --git a/src/backend/common/services/batch_service.py b/src/backend/common/services/batch_service.py index 27a874e0..efce27fa 100644 --- a/src/backend/common/services/batch_service.py +++ b/src/backend/common/services/batch_service.py @@ -288,7 +288,7 @@ async def upload_file_to_batch(self, batch_id: str, user_id: str, file: UploadFi self.logger.info("File uploaded to blob storage", filename=file.filename, batch_id=batch_id) # Create file entry - file_record_obj = await self.database.add_file(batch_id, file_id, file.filename, blob_path) + file_record_obj = await self.database.add_file(UUID(batch_id), UUID(file_id), file.filename, blob_path) file_record_dict = getattr(file_record_obj, "dict", None) file_record = file_record_dict() if callable(file_record_dict) else file_record_obj diff --git a/src/tests/backend/common/database/cosmosdb_test.py b/src/tests/backend/common/database/cosmosdb_test.py index d08fb015..faaaaa09 100644 --- a/src/tests/backend/common/database/cosmosdb_test.py +++ b/src/tests/backend/common/database/cosmosdb_test.py @@ -11,7 +11,7 @@ from uuid import uuid4 # noqa: E402 from azure.cosmos.aio import CosmosClient # noqa: E402 -from azure.cosmos.exceptions import CosmosResourceExistsError # noqa: E402 +from azure.cosmos.exceptions import CosmosResourceExistsError, CosmosResourceNotFoundError # noqa: E402 from common.database.cosmosdb import ( # noqa: E402 CosmosDBClient, @@ -188,6 +188,84 @@ async def test_create_batch_exists(cosmos_db_client, mocker): ) +@pytest.mark.asyncio +async def test_create_batch_conflict_retry_on_404(cosmos_db_client, mocker): + """Test that read_item is retried when it returns 404 after a 409 conflict.""" + user_id = "user_1" + batch_id = uuid4() + + mock_batch_container = mock.MagicMock() + mocker.patch.object(cosmos_db_client, 'batch_container', mock_batch_container) + mock_batch_container.create_item = AsyncMock(side_effect=CosmosResourceExistsError) + + existing_batch = { + "id": str(batch_id), + "batch_id": str(batch_id), + "user_id": user_id, + "file_count": 0, + "created_at": datetime.now(timezone.utc).isoformat(), + "updated_at": datetime.now(timezone.utc).isoformat(), + "status": ProcessStatus.READY_TO_PROCESS, + } + + # First call raises 404, second call succeeds + mock_batch_container.read_item = AsyncMock( + side_effect=[CosmosResourceNotFoundError(message="Not found"), existing_batch] + ) + + batch = await cosmos_db_client.create_batch(user_id, batch_id) + + assert batch.batch_id == batch_id + assert batch.user_id == user_id + assert mock_batch_container.read_item.call_count == 2 + + +@pytest.mark.asyncio +async def test_create_batch_conflict_cross_user(cosmos_db_client, mocker): + """Test that a PermissionError is raised when the batch belongs to a different user.""" + user_id = "user_1" + batch_id = uuid4() + + mock_batch_container = mock.MagicMock() + mocker.patch.object(cosmos_db_client, 'batch_container', mock_batch_container) + mock_batch_container.create_item = AsyncMock(side_effect=CosmosResourceExistsError) + + existing_batch = { + "id": str(batch_id), + "batch_id": str(batch_id), + "user_id": "different_user", + "file_count": 0, + "created_at": datetime.now(timezone.utc).isoformat(), + "updated_at": datetime.now(timezone.utc).isoformat(), + "status": ProcessStatus.READY_TO_PROCESS, + } + mock_batch_container.read_item = AsyncMock(return_value=existing_batch) + + with pytest.raises(PermissionError, match="Batch not found"): + await cosmos_db_client.create_batch(user_id, batch_id) + + +@pytest.mark.asyncio +async def test_create_batch_conflict_exhausted_retries(cosmos_db_client, mocker): + """Test that RuntimeError is raised when read_item returns 404 after all retries.""" + user_id = "user_1" + batch_id = uuid4() + + mock_batch_container = mock.MagicMock() + mocker.patch.object(cosmos_db_client, 'batch_container', mock_batch_container) + mock_batch_container.create_item = AsyncMock(side_effect=CosmosResourceExistsError) + + # All 3 attempts raise 404 + mock_batch_container.read_item = AsyncMock( + side_effect=CosmosResourceNotFoundError(message="Not found") + ) + + with pytest.raises(RuntimeError, match="already exists but could not be read after retries"): + await cosmos_db_client.create_batch(user_id, batch_id) + + assert mock_batch_container.read_item.call_count == 3 + + @pytest.mark.asyncio async def test_create_batch_exception(cosmos_db_client, mocker): user_id = "user_1" @@ -487,6 +565,8 @@ async def mock_query_items(query, parameters, **kwargs): assert file["status"] == ProcessStatus.READY_TO_PROCESS mock_file_container.query_items.assert_called_once() + call_kwargs = mock_file_container.query_items.call_args + assert call_kwargs.kwargs.get("partition_key") == file_id @pytest.mark.asyncio @@ -612,6 +692,8 @@ async def mock_query_items(query, parameters, **kwargs): assert batch["status"] == ProcessStatus.READY_TO_PROCESS mock_batch_container.query_items.assert_called_once() + call_kwargs = mock_batch_container.query_items.call_args + assert call_kwargs.kwargs.get("partition_key") == batch_id @pytest.mark.asyncio