Added more tests

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -30,7 +30,8 @@
 from urllib3.util.retry import Retry
 
 from azure.core import PipelineClient
-from ._session_token_helpers import is_compound_session_token, merge_session_tokens
+from ._session_token_helpers import is_compound_session_token, merge_session_tokens, split_compound_session_tokens, \
+    merge_session_tokens_with_same_pkrangeid
 from ._vector_session_token import VectorSessionToken
 from azure.core.credentials import TokenCredential
 from azure.core.paging import ItemPaged
@@ -3337,15 +3338,17 @@ def _get_partition_key_definition(self, collection_link: str) -> Optional[Dict[s
         return partition_key_definition
 
     def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_range):
-        target_feed_range_normalized = target_feed_range.get_normalized_range()
+        target_feed_range_normalized = target_feed_range._feed_range_internal.get_normalized_range()
         # filter out tuples that overlap with target_feed_range and normalizes all the ranges
-        overlapping_ranges = [(feed_range[0].get_normalized_range(), feed_range[1]) for feed_range in feed_ranges_to_session_tokens if
-                              Range.overlaps(target_feed_range_normalized, feed_range[0].get_normalized_range())]
+        overlapping_ranges = [(feed_range_to_session_token[0]._feed_range_internal.get_normalized_range(), feed_range_to_session_token[1])
+                              for feed_range_to_session_token in feed_ranges_to_session_tokens if Range.overlaps(
+                                target_feed_range_normalized, feed_range_to_session_token[0]._feed_range_internal.get_normalized_range())]
         # Is there a feed_range that is a superset of some of the other feed_ranges excluding tuples
         # with compound session tokens?
-        if overlapping_ranges == 0:
+        if len(overlapping_ranges) == 0:
             raise ValueError('There were no overlapping feed ranges with the target.')
 
+        # merge any session tokens that are the same exact feed range
         i = 0
         j = 1
         while i < len(overlapping_ranges) and j < len(overlapping_ranges):
@@ -3368,7 +3371,6 @@ def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_
 
 
         updated_session_token = ""
-        remaining_session_tokens = []
         done_overlapping_ranges = []
         while len(overlapping_ranges) != 0:
             feed_range_cmp, session_token_cmp = overlapping_ranges[0]
@@ -3393,75 +3395,56 @@ def _get_updated_session_token(self, feed_ranges_to_session_tokens, target_feed_
                 merged_range = subsets[j][0]
                 session_tokens = [subsets[j][1]]
                 merged_indices = [subsets[j][2]]
-                for k in range(len(subsets)):
-                    if j == k:
-                        continue
-                    if merged_range.can_merge(subsets[k][0]):
-                        merged_range = merged_range.merge(subsets[k][0])
-                        session_tokens.append(subsets[k][1])
-                        merged_indices.append(subsets[k][2])
-                    if feed_range_cmp == merged_range:
-                        # if it is the bigger one remove the smaller ranges
-                        # if it is the smaller ranges remove the bigger range
-                        # if it is neither compound
-                        child_lsns_larger = True
-                        child_lsns_smaller = True
-                        for session_token in session_tokens:
-                            tokens = session_token.split(":")
-                            vector_session_token = VectorSessionToken.create(tokens[1])
-                            if vector_session_token.global_lsn <  vector_session_token_cmp.global_lsn:
-                                child_lsns_smaller = False
-                            else:
-                                child_lsns_larger = False
-                        feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices]
-                        for feed_range_to_remove in feed_ranges_to_remove:
-                            overlapping_ranges.remove(feed_range_to_remove)
-                        if child_lsns_larger:
-                            session_tokens.remove(session_token_cmp)
-                            overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens))))
-                            overlapping_ranges.remove(overlapping_ranges[0])
-                        elif child_lsns_smaller:
-                            overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens))))
-                            overlapping_ranges.remove(overlapping_ranges[0])
-
-                        not_found = False
-                        break
+                if len(subsets) == 1:
+                    tokens = session_tokens[0].split(":")
+                    vector_session_token = VectorSessionToken.create(tokens[1])
+                    if vector_session_token_cmp.is_greater(vector_session_token):
+                        overlapping_ranges.remove(overlapping_ranges[merged_indices[0]])
+                else:
+                    for k in range(len(subsets)):
+                        if j == k:
+                            continue
+                        if merged_range.can_merge(subsets[k][0]):
+                            merged_range = merged_range.merge(subsets[k][0])
+                            session_tokens.append(subsets[k][1])
+                            merged_indices.append(subsets[k][2])
+                        if feed_range_cmp == merged_range:
+                            # if it is the bigger one remove the smaller ranges
+                            # if it is the smaller ranges remove the bigger range
+                            # if it is neither compound
+                            child_lsns_larger = True
+                            for session_token in session_tokens:
+                                tokens = session_token.split(":")
+                                vector_session_token = VectorSessionToken.create(tokens[1])
+                                if vector_session_token_cmp.is_greater(vector_session_token):
+                                    child_lsns_larger = False
+                            feed_ranges_to_remove = [overlapping_ranges[i] for i in merged_indices]
+                            for feed_range_to_remove in feed_ranges_to_remove:
+                                overlapping_ranges.remove(feed_range_to_remove)
+                            if child_lsns_larger:
+                                overlapping_ranges.append((merged_range, ','.join(map(str, session_tokens))))
+                                overlapping_ranges.remove(overlapping_ranges[0])
+                            not_found = False
+                            break
+
                 j += 1
 
             done_overlapping_ranges.append(overlapping_ranges[0])
             overlapping_ranges.remove(overlapping_ranges[0])
 
-        for _, session_token in done_overlapping_ranges:
-            # here break up session tokens that are compound
-            if is_compound_session_token(session_token):
-                tokens = session_token.split(",")
-                for token in tokens:
-                    remaining_session_tokens.append(token)
-            else:
-                remaining_session_tokens.append(session_token)
+        # break up session tokens that are compound
+        remaining_session_tokens = split_compound_session_tokens(done_overlapping_ranges)
 
         if len(remaining_session_tokens) == 1:
             return remaining_session_tokens[0]
-        new_session_tokens = []
         # merging any session tokens with same pkrangeid
+        remaining_session_tokens = merge_session_tokens_with_same_pkrangeid(remaining_session_tokens)
+
+        # compound the remaining session tokens
         for i in range(len(remaining_session_tokens)):
-            for j in range(i + 1, len(remaining_session_tokens)):
-                tokens1 = remaining_session_tokens[i].split(":")
-                tokens2 = remaining_session_tokens[j].split(":")
-                pk_range_id1 = tokens1[0]
-                pk_range_id2 = tokens2[0]
-                if pk_range_id1 == pk_range_id2:
-                    vector_session_token1 = VectorSessionToken.create(tokens1[1])
-                    vector_session_token2 = VectorSessionToken.create(tokens2[1])
-                    vector_session_token = vector_session_token1.merge(vector_session_token2)
-                    new_session_tokens.append(pk_range_id1 + ":" + vector_session_token.session_token)
-                    remaining_session_tokens.remove(remaining_session_tokens[i])
-                    remaining_session_tokens.remove(remaining_session_tokens[j])
-        new_session_tokens.extend(remaining_session_tokens)
-        for i in range(len(new_session_tokens)):
-            if i == len(new_session_tokens) - 1:
-                updated_session_token += new_session_tokens[i]
+            if i == len(remaining_session_tokens) - 1:
+                updated_session_token += remaining_session_tokens[i]
             else:
-                updated_session_token += new_session_tokens[i] + ","
+                updated_session_token += remaining_session_tokens[i] + ","
 
         return updated_session_token

sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_range.py

@@ -186,7 +186,6 @@ def _compare_helper(a, b):
 
     @staticmethod
     def overlaps(range1, range2):
-
         if range1 is None or range2 is None:
             return False
         if range1.isEmpty() or range2.isEmpty():
@@ -195,7 +194,7 @@ def overlaps(range1, range2):
         cmp1 = Range._compare_helper(range1.min, range2.max)
         cmp2 = Range._compare_helper(range2.min, range1.max)
 
-        if cmp1 <= 0 or cmp2 <= 0:
+        if cmp1 <= 0 and cmp2 <= 0:
             if (cmp1 == 0 and not (range1.isMinInclusive and range2.isMaxInclusive)) or (
                 cmp2 == 0 and not (range2.isMinInclusive and range1.isMaxInclusive)
             ):
@@ -225,4 +224,4 @@ def is_subset(self, parent_range) -> bool:
         normalized_child_range = self.to_normalized_range()
         return normalized_parent_range.contains(normalized_child_range.min) and \
             (normalized_parent_range.contains(normalized_child_range.max)
-             or normalized_parent_range.max == normalized_child_range.max)
+             or normalized_parent_range.max == normalized_child_range.max)

sdk/cosmos/azure-cosmos/azure/cosmos/_session_token_helpers.py

@@ -36,12 +36,45 @@ def merge_session_tokens(session_token1, session_token2):
     pk_range_id = pk_range_id1
     if pk_range_id1 != pk_range_id2:
         pk_range_id = pk_range_id1 \
-            if vector_session_token1.global_lsn > vector_session_token2.global_lsn else pk_range_id2
+            if vector_session_token1.is_greater(vector_session_token2) else pk_range_id2
     vector_session_token = vector_session_token1.merge(vector_session_token2)
     return pk_range_id + ":" +  vector_session_token.session_token
 
 def is_compound_session_token(session_token):
     return "," in session_token
 
+def split_compound_session_tokens(compound_session_tokens):
+    session_tokens = []
+    for _, session_token in compound_session_tokens:
+            if is_compound_session_token(session_token):
+                tokens = session_token.split(",")
+                for token in tokens:
+                    session_tokens.append(token)
+            else:
+                session_tokens.append(session_token)
+    return session_tokens
 
+def merge_session_tokens_with_same_pkrangeid(session_tokens):
+    new_session_tokens = []
+    i = 0
+    while i < len(session_tokens):
+        j = i + 1
+        while j < len(session_tokens):
+            tokens1 = session_tokens[i].split(":")
+            tokens2 = session_tokens[j].split(":")
+            pk_range_id1 = tokens1[0]
+            pk_range_id2 = tokens2[0]
+            if pk_range_id1 == pk_range_id2:
+                vector_session_token1 = VectorSessionToken.create(tokens1[1])
+                vector_session_token2 = VectorSessionToken.create(tokens2[1])
+                vector_session_token = vector_session_token1.merge(vector_session_token2)
+                new_session_tokens.append(pk_range_id1 + ":" + vector_session_token.session_token)
+                remove_session_tokens = [session_tokens[i], session_tokens[j]]
+                for token in remove_session_tokens:
+                    session_tokens.remove(token)
+                i = -1
+            j += 1
+        i += 1
 
+    new_session_tokens.extend(session_tokens)
+    return new_session_tokens

sdk/cosmos/azure-cosmos/azure/cosmos/container.py

@@ -36,9 +36,7 @@
     _deserialize_throughput,
     _replace_throughput,
     GenerateGuidId,
-    _set_properties_cache,
-    ParsePaths,
-    TrimBeginningAndEndingSlashes
+    _set_properties_cache
 )
 from ._cosmos_client_connection import CosmosClientConnection
 from ._feed_range import FeedRange, FeedRangeEpk

sdk/cosmos/azure-cosmos/test/test_feed_range.py

@@ -11,9 +11,6 @@
 import test_config
 from azure.cosmos._feed_range import FeedRangeEpk
 from azure.cosmos._routing.routing_range import Range
-from test.test_config import TestConfig
-
-
 
 @pytest.fixture(scope="class")
 def setup():
@@ -23,13 +20,58 @@ def setup():
             "You must specify your Azure Cosmos account values for "
             "'masterKey' and 'host' at the top of this class to run the "
             "tests.")
-    test_client = cosmos_client.CosmosClient(TestFeedRange.host, TestConfig.masterKey),
+    test_client = cosmos_client.CosmosClient(TestFeedRange.host, test_config.TestConfig.masterKey),
     created_db = test_client[0].get_database_client(TestFeedRange.TEST_DATABASE_ID)
     return {
         "created_db": created_db,
         "created_collection": created_db.get_container_client(TestFeedRange.TEST_CONTAINER_ID)
     }
 
+test_subset_ranges = [(Range("", "FF", True, False),
+                       Range("3F", "7F", True, False),
+                       True),
+                      (Range("3F", "7F", True, False),
+                Range("", "FF", True, False),
+                False),
+                      (Range("3F", "7F", True, False),
+                Range("", "5F", True, False),
+                False),
+                      (Range("3F", "7F", True, True),
+                Range("3F", "7F", True, True),
+                True),
+                      (Range("3F", "7F", False, True),
+                Range("3F", "7F", True, True),
+                False),
+                      (Range("3F", "7F", True, False),
+                Range("3F", "7F", True, True),
+                False),
+                      (Range("3F", "7F", True, False),
+                Range("", "2F", True, False),
+                False)]
+
+
+test_overlaps_ranges = [(Range("", "FF", True, False),
+                       Range("3F", "7F", True, False),
+                       True),
+                      (Range("3F", "7F", True, False),
+                       Range("", "FF", True, False),
+                       True),
+                      (Range("3F", "7F", True, False),
+                       Range("", "5F", True, False),
+                       True),
+                      (Range("3F", "7F", True, False),
+                       Range("3F", "7F", True, False),
+                       True),
+                      (Range("3F", "7F", True, False),
+                       Range("", "2F", True, False),
+                       False),
+                      (Range("3F", "7F", True, False),
+                       Range("6F", "FF", True, False),
+                       True),
+                      (Range("AA", "BB", True, False),
+                       Range("CC", "FF", True, False),
+                       False)]
+
 @pytest.mark.cosmosEmulator
 @pytest.mark.unittest
 @pytest.mark.usefixtures("setup")
@@ -53,29 +95,7 @@ def test_partition_key_to_feed_range(self, setup):
                         "3c80b1b7310bb39f29cc4ea05bdd461f", True, False)
         setup["created_db"].delete_container(created_container)
 
-    test_ranges = [(Range("", "FF", True, False),
-                    Range("3F", "7F", True, False),
-                    True),
-                   (Range("3F", "7F", True, False),
-                    Range("", "FF", True, False),
-                    False),
-                   (Range("3F", "7F", True, False),
-                    Range("", "5F", True, False),
-                    False),
-                   (Range("3F", "7F", True, True),
-                    Range("3F", "7F", True, True),
-                    True),
-                   (Range("3F", "7F", False, True),
-                    Range("3F", "7F", True, True),
-                    False),
-                   (Range("3F", "7F", True, False),
-                    Range("3F", "7F", True, True),
-                    False),
-                   (Range("3F", "7F", True, False),
-                   Range("", "2F", True, False),
-                   False)]
-
-    @pytest.mark.parametrize("parent_feed_range, child_feed_range, is_subset", test_ranges)
+    @pytest.mark.parametrize("parent_feed_range, child_feed_range, is_subset", test_subset_ranges)
     def test_feed_range_is_subset(self, setup, parent_feed_range, child_feed_range, is_subset):
         epk_parent_feed_range = FeedRangeEpk(parent_feed_range)
         epk_child_feed_range = FeedRangeEpk(child_feed_range)
@@ -86,5 +106,10 @@ def test_feed_range_is_subset_from_pk(self, setup):
         epk_child_feed_range = setup["created_collection"].feed_range_from_partition_key("1")
         assert setup["created_collection"].is_feed_range_subset(epk_parent_feed_range, epk_child_feed_range)
 
+    @pytest.mark.parametrize("range1, range2, overlaps", test_overlaps_ranges)
+    def test_overlaps(self, setup, range1, range2, overlaps):
+        assert Range.overlaps(range1, range2) == overlaps
+
+
 if __name__ == '__main__':
     unittest.main()

sdk/cosmos/azure-cosmos/test/test_request_context.py

@@ -9,9 +9,6 @@
 import azure.cosmos.cosmos_client as cosmos_client
 import test_config
 from azure.cosmos._routing.routing_range import Range
-from test.test_config import TestConfig
-
-
 
 @pytest.fixture(scope="class")
 def setup():
@@ -21,7 +18,7 @@ def setup():
             "You must specify your Azure Cosmos account values for "
             "'masterKey' and 'host' at the top of this class to run the "
             "tests.")
-    test_client = cosmos_client.CosmosClient(TestRequestContext.host, TestConfig.masterKey),
+    test_client = cosmos_client.CosmosClient(TestRequestContext.host, test_config.TestConfig.masterKey),
     created_db = test_client[0].get_database_client(TestRequestContext.TEST_DATABASE_ID)
     return {
         "created_db": created_db,
@@ -52,7 +49,7 @@ def createItem(id = 'item' + str(uuid.uuid4()), pk='A', name='sample'):
 @pytest.mark.unittest
 @pytest.mark.usefixtures("setup")
 class TestRequestContext:
-    """Tests to verify methods for operations on feed ranges
+    """Tests to verify request context gets populated correctly
     """
 
     host = test_config.TestConfig.host