# The MIT License (MIT)
# Copyright (c) Microsoft Corporation. All rights reserved.
import asyncio
import time
import unittest
import random
import uuid
from unittest.mock import patch

import pytest

import test_config
from azure.cosmos import PartitionKey
from azure.cosmos import http_constants
from azure.cosmos import _base
from azure.cosmos.aio import CosmosClient, DatabaseProxy, ContainerProxy
from azure.cosmos.exceptions import CosmosHttpResponseError

async def run_queries(container, iterations):
    ret_list = []
    for i in range(iterations):
        curr = str(random.randint(0, 10))
        query = 'SELECT * FROM c WHERE c.attr1=' + curr + ' order by c.attr1'
        qlist = [item async for item in container.query_items(query=query)]
        ret_list.append((curr, qlist))
    for ret in ret_list:
        curr = ret[0]
        if len(ret[1]) != 0:
            for results in ret[1]:
                attr_number = results['attr1']
                assert str(attr_number) == curr  # verify that all results match their randomly generated attributes
        print("validation succeeded for all query results")


@pytest.mark.cosmosSplit
@pytest.mark.cosmosAADSplit
class TestPartitionSplitQueryAsync(unittest.IsolatedAsyncioTestCase):
    # AAD client/database  -  data-plane (create_item, query_items, _routing_map_provider introspection)
    database: DatabaseProxy = None
    container: ContainerProxy = None
    client: CosmosClient = None
    # Key-auth client/database  -  control-plane (create_container, delete_container, replace_throughput,
    # get_throughput, read_offer)
    key_client: CosmosClient = None
    key_database: DatabaseProxy = None
    key_container: ContainerProxy = None
    configs = test_config.TestConfig
    host = configs.host
    masterKey = configs.masterKey
    throughput = 400
    TEST_DATABASE_ID = configs.TEST_DATABASE_ID
    TEST_CONTAINER_ID = "Single-partition-container-without-throughput-async"
    MAX_TIME = 60 * 10  # 10 minutes for the test to complete, should be enough for partition split to complete

    @classmethod
    def setUpClass(cls):
        if (cls.masterKey == '[YOUR_KEY_HERE]' or
                cls.host == '[YOUR_ENDPOINT_HERE]'):
            raise Exception(
                "You must specify your Azure Cosmos account values for "
                "'masterKey' and 'host' at the top of this class to run the "
                "tests.")

    async def asyncSetUp(self):
        # Control-plane: key-auth (container lifecycle + replace_throughput/get_throughput)
        self.key_client = CosmosClient(self.host, self.masterKey)
        await self.key_client.__aenter__()
        self.key_database = self.key_client.get_database_client(self.TEST_DATABASE_ID)
        self.key_container = await self.key_database.create_container(
            id=self.TEST_CONTAINER_ID,
            partition_key=PartitionKey(path="/id"),
            offer_throughput=self.throughput)
        # Data-plane: AAD
        self.client = test_config.TestConfig.create_data_client_async()
        await self.client.__aenter__()
        self.created_database = self.client.get_database_client(self.TEST_DATABASE_ID)
        self.container = self.created_database.get_container_client(self.TEST_CONTAINER_ID)

    async def asyncTearDown(self):
        try:
            await self.key_database.delete_container(self.TEST_CONTAINER_ID)
        except Exception:
            pass  # Container might not exist if test failed early
        await self.client.close()
        await self.key_client.close()

    async def test_partition_split_query_async(self):
        for i in range(100):
            body = test_config.get_test_item()
            await self.container.create_item(body=body)

        start_time = time.time()
        print("created items, changing offer to 11k and starting queries")
        # Control-plane: replace_throughput via key-auth key_container
        await self.key_container.replace_throughput(11000)
        offer_time = time.time()
        print("changed offer to 11k")
        print("--------------------------------")
        print("now starting queries")

        await run_queries(self.container, 100)  # initial check for queries before partition split
        print("initial check succeeded, now reading offer until replacing is done")
        # Control-plane: get_throughput via key-auth key_container
        offer = await self.key_container.get_throughput()
        while True:
            if time.time() - start_time > self.MAX_TIME:  # timeout test at 10 minutes
                self.skipTest("Partition split didn't complete in time.")
            if offer.properties['content'].get('isOfferReplacePending', False):
                # Keep the event loop responsive while waiting.
                await asyncio.sleep(30)  # wait for the offer to be replaced, check every 30 seconds
                offer = await self.key_container.get_throughput()
            else:
                print("offer replaced successfully, took around {} seconds".format(time.time() - offer_time))
                await run_queries(self.container, 100)  # check queries work post partition split
                self.assertTrue(offer.offer_throughput > self.throughput)
                return

    async def test_incremental_merge_preserves_stable_partitions_async(self):
        """
        Validates the incremental merge logic when a SINGLE partition splits.

        Specifically tests the scenario where:
        - Initial state: 1 partition (< 10k RU/s)
        - After split: 2 partitions (1 parent splits into 2 children)
        - Validation: Both new partitions have parent references

        This test ensures that when ALL partitions split, the incremental merge
        correctly handles the transition without any stable partitions to preserve.
        """
        #Insert data
        for i in range(50):
            item = {'id': f'item_{i}', 'pk': f'pk_{i % 2}', 'value': i}
            await self.container.create_item(item)

        # Verify: We start with 1 partition
        initial_pk_ranges = [r async for r in self.container.client_connection._ReadPartitionKeyRanges(
            self.container.container_link
        )]
        print(f"Initial partition count: {len(initial_pk_ranges)}")
        assert len(initial_pk_ranges) == 1, \
            f"Expected 1 partition at 400 RU/s, got {len(initial_pk_ranges)}"

        # Force initial routing map cache by running a query
        await run_queries(self.container, 1)

        # Trigger split via the shared bounded helper (timeout + SkipTest)
        # instead of an unbounded polling loop.
        await test_config.TestConfig.trigger_split_async(self.key_container, 11000)

        # Run queries to trigger routing map refresh
        await run_queries(self.container, 1)

        # Verify: Split created 2 partitions using _ReadPartitionKeyRanges
        post_split_pk_ranges = [r async for r in self.container.client_connection._ReadPartitionKeyRanges(
            self.container.container_link
        )]
        print(f"Post-split partition count: {len(post_split_pk_ranges)}")

        assert len(post_split_pk_ranges) == 2, \
            f"Expected 2 partitions after split, got {len(post_split_pk_ranges)}"

        # Verify: Access routing map
        provider = self.container.client_connection._routing_map_provider
        # The routing map is cached by collection link
        cached_map = provider._collection_routing_map_by_item.get(
            self.container.container_link
        )

        current_ranges = list(cached_map._orderedPartitionKeyRanges)

        # When a single partition splits, both children will have 'parents'
        child_partitions = [r for r in current_ranges if r.get('parents')]
        stable_partitions = [r for r in current_ranges if not r.get('parents')]

        # This validates ALL partitions split (both have parent references)
        assert len(child_partitions) == 2, \
            f"Expected 2 child partitions with parents, got {len(child_partitions)}"
        assert len(stable_partitions) == 0, \
            f"Expected 0 stable partitions (all split), got {len(stable_partitions)}"

        print(f"Validated: Single partition split into {len(child_partitions)} children")

        # Verify final throughput  -  control-plane
        final_offer = await self.key_container.get_throughput()
        assert final_offer.offer_throughput == 11000

    async def test_incremental_merge_handles_split_partitions_async(self):
        """
        Validates the incremental merge logic when SOME partitions split.

        Specifically tests BOTH branches in _get_routing_map_with_change_feed:
        - Initial state: 2 partitions (11k RU/s)
        - After split: 3 partitions (1 stable + 2 new from split)
        - Validation:
          1. `if parents:` branch - New child partitions inherit parent's range_info
          2. `else:` branch - Stable partition preserves its original range_info

        This test ensures that during a PARTIAL split, the incremental merge correctly:
        - Handles new child partitions (with parent references)
        - Preserves unchanged partitions (without parent references)
        """
        new_container_id = 'partial_split_test_' + str(uuid.uuid4())
        # Control-plane: create via key-auth key_database
        new_setup_container = await self.key_database.create_container(
            id=new_container_id,
            partition_key=PartitionKey(path="/pk"),
            offer_throughput=11000  # 2 physical partitions
        )
        # Data-plane: re-bind via AAD database
        new_container = self.created_database.get_container_client(new_container_id)
        try:
            # Insert data
            for i in range(200):
                item = {'id': f'item_{i}', 'pk': f'pk_{i % 2}', 'value': i}
                await new_container.create_item(item)

            # Verify: We start with 2 partitions
            initial_pk_ranges = [r async for r in new_container.client_connection._ReadPartitionKeyRanges(
                new_container.container_link
            )]
            print(f"Initial partition count: {len(initial_pk_ranges)}")
            assert len(initial_pk_ranges) == 2, \
                f"Expected 2 partitions at 11k RU/s, got {len(initial_pk_ranges)}"

            # Force initial routing map cache
            await run_queries(new_container, 1)

            # Trigger split via the shared bounded helper (timeout + SkipTest)
            # instead of an unbounded polling loop.
            await test_config.TestConfig.trigger_split_async(new_setup_container, 25000)

            # Run queries to trigger routing map refresh
            await run_queries(new_container, 1)

            # Verify: Split created 3 partitions
            post_split_pk_ranges = [r async for r in new_container.client_connection._ReadPartitionKeyRanges(
                new_container.container_link
            )]
            print(f"Post-split partition count: {len(post_split_pk_ranges)}")

            assert len(post_split_pk_ranges) == 3, \
                f"Expected 3 partitions after partial split, got {len(post_split_pk_ranges)}"

            # Verify: Access routing map and validate both code paths
            provider = new_container.client_connection._routing_map_provider
            cached_map = provider._collection_routing_map_by_item.get(
                new_container.container_link
            )

            current_ranges = list(cached_map._orderedPartitionKeyRanges)
            child_partitions = [r for r in current_ranges if r.get('parents')]
            stable_partitions = [r for r in current_ranges if not r.get('parents')]

            # This validates BOTH code paths
            # Test the `if parents:` branch
            assert len(child_partitions) == 2, \
                f"Expected 2 child partitions from split, got {len(child_partitions)}"

            # Test the `else:` branch
            assert len(stable_partitions) == 1, \
                f"Expected 1 stable partition (unchanged), got {len(stable_partitions)}"

            # Verify parent references on children
            for child in child_partitions:
                parents = child.get('parents', [])
                assert len(parents) > 0, f"Child {child['id']} has empty parents"
                print(f"Child partition {child['id']} has parent {parents[0]}")

            print(f"Validated: {len(stable_partitions)} stable + {len(child_partitions)} split partitions")

            # Verify final throughput  -  control-plane
            final_offer = await new_setup_container.get_throughput()
            assert final_offer.offer_throughput == 25000

        finally:
            await self.key_database.delete_container(new_container_id)

    async def test_incremental_change_feed_only_affects_target_collection_async(self):
        """
        Validates that incremental change feed updates are collection-scoped.

        Scenario:
        - Create 2 containers: container_A and container_B
        - Both start with 1 partition (400 RU/s)
        - Run queries on both to populate routing map cache
        - Split ONLY container_A (400 -> 11000 RU/s)
        - Verify:
          1. container_A's routing map is refreshed (2 partitions)
          2. container_B's routing map is unchanged (1 partition)
          3. container_B's cache is NOT invalidated
        """
        container_a_id = 'container_a_async_' + str(uuid.uuid4())
        container_b_id = 'container_b_async_' + str(uuid.uuid4())
        # Control-plane: create via key-auth key_database
        key_container_a = await self.key_database.create_container(
            id=container_a_id,
            partition_key=PartitionKey(path="/pk"),
            offer_throughput=400
        )
        await self.key_database.create_container(
            id=container_b_id,
            partition_key=PartitionKey(path="/pk"),
            offer_throughput=400
        )
        # Data-plane: re-bind via AAD database
        container_a = self.created_database.get_container_client(container_a_id)
        container_b = self.created_database.get_container_client(container_b_id)

        try:
            # Insert data into both containers
            for i in range(50):
                await container_a.create_item({'id': f'a_{i}', 'pk': f'pk_{i}', 'value': i})
                await container_b.create_item({'id': f'b_{i}', 'pk': f'pk_{i}', 'value': i})

            # Force routing map cache population for BOTH containers
            await run_queries(container_a, 1)
            await run_queries(container_b, 1)

            # Verify initial state: Both have 1 partition
            provider = container_a.client_connection._routing_map_provider
            collection_id_a = _base.GetResourceIdOrFullNameFromLink(container_a.container_link)
            collection_id_b = _base.GetResourceIdOrFullNameFromLink(container_b.container_link)

            map_a_before = provider._collection_routing_map_by_item.get(collection_id_a)
            map_b_before = provider._collection_routing_map_by_item.get(collection_id_b)

            assert map_a_before is not None, "Container A should have cached routing map"
            assert map_b_before is not None, "Container B should have cached routing map"

            ranges_a_before = list(map_a_before._orderedPartitionKeyRanges)
            ranges_b_before = list(map_b_before._orderedPartitionKeyRanges)

            assert len(ranges_a_before) == 1, f"Expected 1 partition in A, got {len(ranges_a_before)}"
            assert len(ranges_b_before) == 1, f"Expected 1 partition in B, got {len(ranges_b_before)}"

            # Store container B's original routing map object reference
            map_b_object_id = id(map_b_before)

            print(f"Before split - Container A: {len(ranges_a_before)} partitions")
            print(f"Before split - Container B: {len(ranges_b_before)} partitions")
            print(f"Container B routing map object ID: {map_b_object_id}")

            # Split only Container A via the shared bounded helper.
            await test_config.TestConfig.trigger_split_async(key_container_a, 11000)

            # Wait for physical partition ranges to reflect the split.
            split_convergence_deadline = time.time() + 300
            physical_count = 0
            while time.time() < split_convergence_deadline:
                physical_ranges = [r async for r in container_a.client_connection._ReadPartitionKeyRanges(
                    container_a.container_link
                )]
                physical_count = len(physical_ranges)
                if physical_count == 2:
                    break
                await asyncio.sleep(5)

            assert physical_count == 2, \
                f"Container A physical partitions did not converge to 2 after split, got {physical_count}"

            # Trigger and wait for routing-map cache refresh for container A.
            cached_count = 0
            current_map = map_a_before
            while time.time() < split_convergence_deadline:
                provider = container_a.client_connection._routing_map_provider
                refreshed_map = await provider.get_routing_map(
                    collection_link=container_a.container_link,
                    feed_options={},
                    force_refresh=True,
                    previous_routing_map=current_map,
                )
                current_map = refreshed_map or current_map
                cached_count = len(list(current_map._orderedPartitionKeyRanges)) if current_map else 0
                if cached_count == 2:
                    break
                await asyncio.sleep(5)

            # Run queries on container B to verify cache wasn't invalidated.
            await run_queries(container_b, 1)

            # Verify post-split state
            provider = container_a.client_connection._routing_map_provider
            map_a_after = provider._collection_routing_map_by_item.get(collection_id_a)
            map_b_after = provider._collection_routing_map_by_item.get(collection_id_b)

            assert map_a_after is not None, "Container A should still have cached routing map"
            assert map_b_after is not None, "Container B should still have cached routing map"

            ranges_a_after = list(map_a_after._orderedPartitionKeyRanges)
            ranges_b_after = list(map_b_after._orderedPartitionKeyRanges)

            print(f"After split - Container A: {len(ranges_a_after)} partitions")
            print(f"After split - Container B: {len(ranges_b_after)} partitions")

            # Validate container A was updated
            assert len(ranges_a_after) == 2, \
                f"Container A should have 2 partitions after split, got {len(ranges_a_after)} (physical={physical_count}, cached={cached_count})"

            # Validate container B remained unchanged
            assert len(ranges_b_after) == 1, \
                f"Container B should still have 1 partition, got {len(ranges_b_after)}"

            # Verify container B's cache wasn't cleared/recreated
            # The routing map object should be the SAME instance
            map_b_after_object_id = id(map_b_after)
            assert map_b_object_id == map_b_after_object_id, \
                f"Container B's routing map was recreated (cache invalidated). " \
                f"Before: {map_b_object_id}, After: {map_b_after_object_id}"

            # Verify the partition IDs haven't changed for container B
            assert ranges_b_before[0]['id'] == ranges_b_after[0]['id'], \
                "Container B's partition ID changed unexpectedly"

            print("Validated: Incremental change feed only affected container A")
            print("Container B's routing map remained untouched (same object reference)")

        finally:
            await self.key_database.delete_container(container_a_id)
            await self.key_database.delete_container(container_b_id)

    async def test_routing_map_provider_fallback_on_incomplete_merge_async(self):
        """
        Validates that routing_map_provider falls back to full refresh
        when incremental merge produces incomplete range coverage.
        """
        container_id = 'test_fallback_async_' + str(uuid.uuid4())
        await self.key_database.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/pk"),
            offer_throughput=400
        )
        container = self.created_database.get_container_client(container_id)

        try:
            # Insert data
            for i in range(20):
                await container.create_item({'id': f'item_{i}', 'pk': f'pk_{i}', 'value': i})

            # Populate cache
            await run_queries(container, 1)

            provider = container.client_connection._routing_map_provider
            collection_link = container.container_link

            # Get initial routing map
            initial_map = await provider.get_routing_map(
                collection_link=collection_link,
                feed_options={}
            )

            assert initial_map is not None
            initial_ranges = list(initial_map._orderedPartitionKeyRanges)
            assert len(initial_ranges) == 1

            # Simulate a corrupted incremental update by creating incomplete ranges
            parent_id = initial_ranges[0]['id']

            # Create child ranges with a GAP (incomplete coverage)
            corrupted_child_1 = {
                'id': 'corrupted_1',
                'minInclusive': '',
                'maxExclusive': '80',
                'parents': [parent_id]
            }

            # Gap from '80' to '90' - missing coverage!
            corrupted_child_2 = {
                'id': 'corrupted_2',
                'minInclusive': '90',  # Should be '80'
                'maxExclusive': 'FF',
                'parents': [parent_id]
            }

            # Test try_combine with corrupted data
            combined_map = initial_map.try_combine(
                [(corrupted_child_1, True), (corrupted_child_2, True)],
                new_change_feed_etag="corrupted_etag"
            )

            # Should return None due to incomplete coverage
            assert combined_map is None, \
                "try_combine should return None for incomplete range set"

            print("Validated: try_combine correctly rejects incomplete ranges")

            # Now verify the routing_map_provider's fallback behavior
            # Clear cache to force a fresh fetch
            provider._collection_routing_map_by_item.clear()

            # This should trigger the full refresh fallback path
            refreshed_map = await provider.get_routing_map(
                collection_link=collection_link,
                feed_options={}
            )

            # Should successfully retrieve the routing map via full refresh
            assert refreshed_map is not None, \
                "Routing map should be retrieved via full refresh fallback"

            refreshed_ranges = list(refreshed_map._orderedPartitionKeyRanges)
            assert len(refreshed_ranges) == 1, \
                f"Expected 1 partition after fallback, got {len(refreshed_ranges)}"

            # Verify completeness
            assert refreshed_ranges[0]['minInclusive'] == ''
            assert refreshed_ranges[0]['maxExclusive'] == 'FF'

            print("Validated: routing_map_provider successfully fell back to full refresh")

            # Verify queries still work after fallback
            query_results = [item async for item in container.query_items(
                query='SELECT * FROM c',
            )]

            assert len(query_results) == 20, \
                f"Expected 20 items after fallback, got {len(query_results)}"

            print("Validated: Queries work correctly after fallback")

        finally:
            await self.key_database.delete_container(container_id)

    async def test_is_cache_stale_etag_comparison_async(self):
        """
        Validates _is_cache_stale() ETag logic:
        - Returns True when cached map's ETag matches previous_routing_map's ETag (stale)
        - Returns False when ETags differ (already refreshed by another thread)
        - Returns False when previous_routing_map is None
        - Returns False when cache is empty
        """
        from unittest.mock import MagicMock

        container_id = 'test_stale_etag_async_' + str(uuid.uuid4())
        await self.key_database.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/pk"),
            offer_throughput=400
        )
        container = self.created_database.get_container_client(container_id)

        try:
            for i in range(10):
                await container.create_item({'id': f'item_{i}', 'pk': f'pk_{i}', 'value': i})

            # Populate cache
            await run_queries(container, 1)

            provider = container.client_connection._routing_map_provider
            collection_link = container.container_link
            collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)

            cached_map = provider._collection_routing_map_by_item.get(collection_id)
            assert cached_map is not None

            # Case 1: previous_routing_map is None -> should return False
            assert provider._is_cache_stale(collection_id, None) is False, \
                "_is_cache_stale should return False when previous_routing_map is None"

            # Case 2: ETags match -> should return True (cache is stale)
            assert provider._is_cache_stale(collection_id, cached_map) is True, \
                "_is_cache_stale should return True when ETags match"

            # Case 3: ETags differ -> should return False (already refreshed)
            fake_previous = MagicMock()
            fake_previous.change_feed_etag = "fake-different-etag-12345"
            assert provider._is_cache_stale(collection_id, fake_previous) is False, \
                "_is_cache_stale should return False when ETags differ"

            # Case 4: Cache is empty -> should return False
            provider._collection_routing_map_by_item.clear()
            real_previous = MagicMock()
            real_previous.change_feed_etag = cached_map.change_feed_etag
            assert provider._is_cache_stale(collection_id, real_previous) is False, \
                "_is_cache_stale should return False when cache is empty"

            print("Validated: _is_cache_stale ETag comparison logic works correctly")

        finally:
            await self.key_database.delete_container(container_id)

    async def test_full_load_with_incomplete_ranges_surfaces_503_async(self):
        """
        Validates that a full load with incomplete ranges surfaces a retryable
        HTTP 503 after exhausting the bounded retry budget.
        When a full load is performed (previous_routing_map=None) and the service
        returns gapped ranges, _fetch_routing_map should not leak internal
        map-construction failures to callers.
        """
        container_id = 'test_fallback_guard_async_' + str(uuid.uuid4())
        await self.key_database.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/pk"),
            offer_throughput=400
        )
        container = self.created_database.get_container_client(container_id)

        try:
            for i in range(10):
                await container.create_item({'id': f'item_{i}', 'pk': f'pk_{i}', 'value': i})

            await run_queries(container, 1)

            provider = container.client_connection._routing_map_provider
            collection_link = container.container_link
            collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)

            # Create an incomplete range set that will cause CompleteRoutingMap to fail
            incomplete_range = {
                'id': 'incomplete_0',
                'minInclusive': '',
                'maxExclusive': '80',  # Gap from 80 to FF
            }

            async def mock_read_ranges(*args, **kwargs):
                # Mirror the production wire-up: _asynchronous_request populates
                # this sidecar with the real HTTP status. Without it, the drain
                # loop's status==304 termination contract can't trip and the
                # loop would run unbounded (OOM in CI).
                status_capture = kwargs.get('_internal_response_status_capture')
                if status_capture is not None:
                    status_capture[0] = http_constants.StatusCodes.NOT_MODIFIED
                yield incomplete_range

            with patch.object(
                    container.client_connection,
                    '_ReadPartitionKeyRanges',
                    side_effect=mock_read_ranges
            ):
                async def _no_sleep(_seconds):
                    return None

                with patch('azure.cosmos._routing.aio.routing_map_provider.asyncio.sleep', new=_no_sleep):
                    with self.assertRaises(CosmosHttpResponseError) as ctx:
                        await provider._fetch_routing_map(
                            collection_link=collection_link,
                            collection_id=collection_id,
                            previous_routing_map=None,
                            feed_options={},
                        )
                self.assertEqual(ctx.exception.status_code, http_constants.StatusCodes.SERVICE_UNAVAILABLE)

            print("Validated: full load with incomplete ranges surfaces retryable HTTP 503")

        finally:
            await self.key_database.delete_container(container_id)

    async def test_internal_pk_range_fetch_flag_is_set_async(self):
        """
        Validates that _fetch_routing_map sets the _internal_pk_range_fetch flag
        in the options passed to _ReadPartitionKeyRanges. This flag prevents
        infinite recursion when the PK range fetch itself gets a 410.
        """
        container_id = 'test_pk_flag_async_' + str(uuid.uuid4())
        await self.key_database.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/pk"),
            offer_throughput=400
        )
        container = self.created_database.get_container_client(container_id)

        try:
            for i in range(10):
                await container.create_item({'id': f'item_{i}', 'pk': f'pk_{i}', 'value': i})

            provider = container.client_connection._routing_map_provider
            collection_link = container.container_link

            captured_options = {}
            original_read = container.client_connection._ReadPartitionKeyRanges

            def spy_read_ranges(*args, **kwargs):
                if len(args) > 1:
                    captured_options.update(args[1])
                return original_read(*args, **kwargs)

            with patch.object(
                    container.client_connection,
                    '_ReadPartitionKeyRanges',
                    side_effect=spy_read_ranges
            ):
                # Clear cache to force a fresh fetch
                provider._collection_routing_map_by_item.clear()

                routing_map = await provider.get_routing_map(
                    collection_link=collection_link,
                    feed_options={}
                )

                assert routing_map is not None
                assert captured_options.get("_internal_pk_range_fetch") is True, \
                    "_internal_pk_range_fetch flag should be set in options"

            print("Validated: _internal_pk_range_fetch flag is correctly set")

        finally:
            await self.key_database.delete_container(container_id)

    async def test_lock_free_fast_path_async(self):
        """
        Validates the lock-free fast path in get_routing_map().
        When the cache is populated and no force_refresh is requested,
        the map should be returned without acquiring the collection lock.
        """
        container_id = 'test_fast_path_async_' + str(uuid.uuid4())
        await self.key_database.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/pk"),
            offer_throughput=400
        )
        container = self.created_database.get_container_client(container_id)

        try:
            for i in range(10):
                await container.create_item({'id': f'item_{i}', 'pk': f'pk_{i}', 'value': i})

            # Populate cache
            await run_queries(container, 1)

            provider = container.client_connection._routing_map_provider
            collection_link = container.container_link
            collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)

            # Verify cache is populated
            cached_map = provider._collection_routing_map_by_item.get(collection_id)
            assert cached_map is not None

            # Spy on _get_lock_for_collection to verify it's NOT called on the fast path
            lock_call_count = {'count': 0}
            original_get_lock = provider._get_lock_for_collection

            async def spy_get_lock(*args, **kwargs):
                lock_call_count['count'] += 1
                return await original_get_lock(*args, **kwargs)

            with patch.object(provider, '_get_lock_for_collection', side_effect=spy_get_lock):
                # This should hit the fast path  -  no lock acquisition
                result = await provider.get_routing_map(
                    collection_link=collection_link,
                    feed_options={}
                )

                assert result is cached_map, "Should return the same cached map object"
                assert lock_call_count['count'] == 0, \
                    "Lock should NOT be acquired on the fast path"

            # Now verify that force_refresh DOES acquire the lock
            lock_call_count['count'] = 0
            with patch.object(provider, '_get_lock_for_collection', side_effect=spy_get_lock):
                result = await provider.get_routing_map(
                    collection_link=collection_link,
                    feed_options={},
                    force_refresh=True
                )
                assert result is not None
                assert lock_call_count['count'] == 1, \
                    "Lock SHOULD be acquired when force_refresh=True"

            print("Validated: Lock-free fast path works correctly")

        finally:
            await self.key_database.delete_container(container_id)

    async def test_response_hook_chaining_async(self):
        """
        Validates that an upstream response_hook is preserved and called
        alongside _fetch_routing_map's internal capture_response_hook.
        Both hooks should receive the response headers.
        """
        container_id = 'test_hook_chain_async_' + str(uuid.uuid4())
        await self.key_database.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/pk"),
            offer_throughput=400
        )
        container = self.created_database.get_container_client(container_id)

        try:
            for i in range(10):
                await container.create_item({'id': f'item_{i}', 'pk': f'pk_{i}', 'value': i})

            provider = container.client_connection._routing_map_provider
            collection_link = container.container_link
            collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)

            # Clear cache to force a fetch
            provider._collection_routing_map_by_item.clear()

            upstream_headers_captured = {}

            def upstream_hook(headers, body):
                upstream_headers_captured.update(headers)

            # Call get_routing_map with an upstream response_hook in kwargs
            result = await provider.get_routing_map(
                collection_link=collection_link,
                feed_options={},
                response_hook=upstream_hook
            )

            assert result is not None, "Routing map should be returned"

            # The internal hook should have captured the ETag (used for change_feed_etag)
            assert result.change_feed_etag is not None, \
                "Internal capture hook should have captured the ETag"

            # The upstream hook should also have been called with headers
            assert len(upstream_headers_captured) > 0, \
                "Upstream response_hook should have been called with headers"

            print("Validated: response_hook chaining works correctly")

        finally:
            await self.key_database.delete_container(container_id)

    async def test_if_none_match_header_cleanup_on_fallback_async(self):
        """
        Validates that stale IfNoneMatch headers are removed when falling back
        from incremental to full load.

        Current behavior includes one incremental retry before full refresh.
        This test forces both incremental attempts to be incomplete, then
        verifies the final full-load call drops IfNoneMatch.
        """
        container_id = 'test_etag_cleanup_async_' + str(uuid.uuid4())
        await self.key_database.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/pk"),
            offer_throughput=400
        )
        container = self.created_database.get_container_client(container_id)

        try:
            for i in range(10):
                await container.create_item({'id': f'item_{i}', 'pk': f'pk_{i}', 'value': i})

            # Populate cache
            await run_queries(container, 1)

            provider = container.client_connection._routing_map_provider
            collection_link = container.container_link
            collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)

            # Capture the cached map to use as the stale previous_routing_map
            cached_map = provider._collection_routing_map_by_item.get(collection_id)
            assert cached_map is not None

            captured_headers_list = []
            original_read = container.client_connection._ReadPartitionKeyRanges

            call_count = {'count': 0}

            def spy_read_ranges(*args, **kwargs):
                call_count['count'] += 1
                # Capture the headers from kwargs
                headers = kwargs.get('headers', {})
                captured_headers_list.append(headers.copy())

                if call_count['count'] <= 2:
                    # First two calls are incremental attempts; return a child
                    # with a missing parent so merge is incomplete and fallback
                    # path is exercised. Mirror the production wire-up:
                    # _asynchronous_request populates this sidecar with the real
                    # HTTP status. Without it, the drain loop's status==304
                    # termination contract can't trip and evaluate_drain_page
                    # raises RuntimeError.
                    status_capture = kwargs.get('_internal_response_status_capture')
                    if status_capture is not None:
                        status_capture[0] = http_constants.StatusCodes.NOT_MODIFIED
                    fake_child = {
                        'id': f'child_{call_count["count"]}',
                        'minInclusive': '',
                        'maxExclusive': 'FF',
                        'parents': ['nonexistent_parent']
                    }
                    async def gen():
                        yield fake_child
                    return gen()
                # Third call should be full-load fallback.
                return original_read(*args, **kwargs)

            with patch.object(
                    container.client_connection,
                    '_ReadPartitionKeyRanges',
                    side_effect=spy_read_ranges
            ):
                # Force refresh with the stale map to trigger incremental path.
                result = await provider.get_routing_map(
                    collection_link=collection_link,
                    feed_options={},
                    force_refresh=True,
                    previous_routing_map=cached_map
                )

                assert result is not None

                # Expected call sequence:
                #   1. incremental attempt        (IfNoneMatch = stale etag)
                #   2. incremental retry          (IfNoneMatch = stale etag)
                #   3. full-load fallback page 1  (no IfNoneMatch -- the cleanup we are testing)
                #   4. full-load fallback page 2  (IfNoneMatch = FRESH etag from page 1,
                #                                  to receive the 304 terminator that ends
                #                                  the drain loop -- peer-SDK parity)
                stale_etag = cached_map.change_feed_etag
                assert call_count['count'] >= 3, \
                    f"Expected at least 3 calls to _ReadPartitionKeyRanges, got {call_count['count']}"

                # First two calls should be incremental and carry the stale IfNoneMatch.
                first_headers = captured_headers_list[0]
                assert first_headers.get(http_constants.HttpHeaders.IfNoneMatch) == stale_etag, \
                    "First call (incremental) should have stale IfNoneMatch header"

                second_headers = captured_headers_list[1]
                assert second_headers.get(http_constants.HttpHeaders.IfNoneMatch) == stale_etag, \
                    "Second call (incremental retry) should have stale IfNoneMatch header"

                # Third call is full-load fallback and MUST drop IfNoneMatch -- this is
                # the bug fix's whole point. Any post-fallback drain pages (call 4+)
                # legitimately reuse the etag returned by call 3 as their If-None-Match
                # to receive the 304 terminator; that fresh etag may coincidentally equal
                # the original stale etag if nothing changed server-side between caching
                # and fallback, so we cannot assert "!= stale_etag" on those drain pages.
                # The call-3 assertion is the actual production contract.
                third_headers = captured_headers_list[2]
                assert http_constants.HttpHeaders.IfNoneMatch not in third_headers, \
                    "Third call (full load fallback) should NOT have IfNoneMatch header"

            print("Validated: IfNoneMatch header is correctly cleaned up on fallback")

        finally:
            await self.key_database.delete_container(container_id)

    async def test_force_refresh_with_stale_previous_routing_map_async(self):
        """
        Validates get_routing_map(force_refresh=True, previous_routing_map=stale_map)
        correctly refreshes the cache and queries continue working.
        This tests the targeted refresh path used by the 410 retry policy.
        """
        container_id = 'test_force_refresh_async_' + str(uuid.uuid4())
        await self.key_database.create_container(
            id=container_id,
            partition_key=PartitionKey(path="/pk"),
            offer_throughput=400
        )
        container = self.created_database.get_container_client(container_id)

        try:
            for i in range(20):
                await container.create_item({'id': f'item_{i}', 'pk': f'pk_{i}', 'value': i})

            # Populate cache
            await run_queries(container, 1)

            provider = container.client_connection._routing_map_provider
            collection_link = container.container_link
            collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)

            # Capture the current map as the "stale" previous map
            stale_map = provider._collection_routing_map_by_item.get(collection_id)
            assert stale_map is not None
            original_etag = stale_map.change_feed_etag

            # Force refresh with the stale map  -  simulates what the gone retry policy does
            refreshed_map = await provider.get_routing_map(
                collection_link=collection_link,
                feed_options={},
                force_refresh=True,
                previous_routing_map=stale_map
            )

            assert refreshed_map is not None
            assert len(list(refreshed_map._orderedPartitionKeyRanges)) >= 1

            # Verify queries still work after force refresh
            results = [item async for item in container.query_items(
                query='SELECT * FROM c',
            )]
            assert len(results) == 20, \
                f"Expected 20 items after force refresh, got {len(results)}"

            print("Validated: Force refresh with previous_routing_map works correctly")

        finally:
            await self.key_database.delete_container(container_id)


if __name__ == '__main__':
    unittest.main()


