Add files using upload-large-folder tool

61ba51e verified 7 months ago

7.74 kB

import random
import string
import time
import unittest
from typing import Dict, List, Tuple

from tree import MultiTenantRadixTree


class TestMultiTenantRadixTree(unittest.TestCase):
    def setUp(self):
        self.tree = MultiTenantRadixTree()

    def test_insert_exact_match(self):
        """Test 1: Basic insert and exact match operations"""
        # Insert a single string for one tenant
        self.tree.insert("hello", "tenant1")
        matched, tenant = self.tree.prefix_match("hello")
        self.assertEqual(matched, "hello")
        self.assertEqual(tenant, "tenant1")

        # Insert same string for different tenant
        self.tree.insert("hello", "tenant2")
        matched, tenant = self.tree.prefix_match("hello")
        self.assertIn(tenant, ["tenant1", "tenant2"])

        # Insert different string for same tenant
        self.tree.insert("world", "tenant1")
        matched, tenant = self.tree.prefix_match("world")
        self.assertEqual(matched, "world")
        self.assertEqual(tenant, "tenant1")

        print(self.tree.pretty_print())

    def test_insert_partial_match(self):
        """Test 2: Insert with partial matching scenarios"""
        # Test partial matches with common prefixes
        self.tree.insert("hello", "tenant1")
        print(self.tree.pretty_print())
        self.tree.insert("help", "tenant2")
        print(self.tree.pretty_print())

        # Match exact strings
        matched, tenant = self.tree.prefix_match("hello")
        self.assertEqual(matched, "hello")
        self.assertEqual(tenant, "tenant1")

        matched, tenant = self.tree.prefix_match("help")
        self.assertEqual(matched, "help")
        self.assertEqual(tenant, "tenant2")

        # Match partial string
        matched, tenant = self.tree.prefix_match("hel")
        self.assertEqual(matched, "hel")
        self.assertIn(tenant, ["tenant1", "tenant2"])

        # Match longer string
        matched, tenant = self.tree.prefix_match("hello_world")
        self.assertEqual(matched, "hello")
        self.assertEqual(tenant, "tenant1")

    def test_insert_edge_cases(self):
        """Test 3: Edge cases for insert and match operations"""
        # Empty string
        self.tree.insert("", "tenant1")
        matched, tenant = self.tree.prefix_match("")
        self.assertEqual(matched, "")
        self.assertEqual(tenant, "tenant1")

        # Single character
        self.tree.insert("a", "tenant1")
        matched, tenant = self.tree.prefix_match("a")
        self.assertEqual(matched, "a")
        self.assertEqual(tenant, "tenant1")

        # Very long string
        long_str = "a" * 1000
        self.tree.insert(long_str, "tenant1")
        matched, tenant = self.tree.prefix_match(long_str)
        self.assertEqual(matched, long_str)
        self.assertEqual(tenant, "tenant1")

        # Unicode characters
        self.tree.insert("你好", "tenant1")
        matched, tenant = self.tree.prefix_match("你好")
        self.assertEqual(matched, "你好")
        self.assertEqual(tenant, "tenant1")

    def test_simple_eviction(self):
        """Test 4: Simple eviction scenarios
        Tenant1: limit 10 chars
        Tenant2: limit 5 chars
        Should demonstrate:
        1. Basic eviction when size limit exceeded
        2. Proper eviction based on last access time
        3. Verification that shared nodes remain intact for other tenants
        """
        # Set up size limits
        max_size = {"tenant1": 10, "tenant2": 5}

        # Insert strings for both tenants
        self.tree.insert("hello", "tenant1")  # size 5
        self.tree.insert("hello", "tenant2")  # size 5
        self.tree.insert("world", "tenant2")  # size 5, total for tenant2 = 10

        # Verify initial sizes
        sizes_before = self.tree.get_used_size_per_tenant()
        self.assertEqual(sizes_before["tenant1"], 5)  # "hello" = 5
        self.assertEqual(sizes_before["tenant2"], 10)  # "hello" + "world" = 10

        # Evict - should remove "hello" from tenant2 as it's the oldest
        self.tree.evict_tenant_data(max_size)

        # Verify sizes after eviction
        sizes_after = self.tree.get_used_size_per_tenant()
        self.assertEqual(sizes_after["tenant1"], 5)  # Should be unchanged
        self.assertEqual(sizes_after["tenant2"], 5)  # Only "world" remains

        # Verify "world" remains for tenant2 (was accessed more recently)
        matched, tenant = self.tree.prefix_match("world")
        self.assertEqual(matched, "world")
        self.assertEqual(tenant, "tenant2")

    def test_medium_eviction(self):
        """Test 5: Medium complexity eviction scenarios with shared prefixes
        Tenant1: limit 10 chars
        Tenant2: limit 7 chars (forces one string to be evicted)
        Tree structure after inserts:
        └── 'h' [t1, t2]
            ├── 'i' [t1, t2]      # Oldest for t2
            └── 'e' [t1, t2]
                ├── 'llo' [t1, t2]
                └── 'y' [t2]      # Newest for t2
        Size calculations:
        tenant1: "h"(1) + "i"(1) + "e"(1) + "llo"(3) = 6 chars
        tenant2: "h"(1) + "i"(1) + "e"(1) + "llo"(3) + "y"(1) = 7 chars
        After eviction (tenant2 exceeds limit by 1 char):
        "hi" should be removed from tenant2 as it's the oldest access
        """
        max_size = {
            "tenant1": 10,
            "tenant2": 6,
        }  # tenant2 will need to evict one string

        # Create a tree with overlapping prefixes
        self.tree.insert("hi", "tenant1")
        self.tree.insert("hi", "tenant2")  # OLDEST for t2

        self.tree.insert("hello", "tenant1")
        self.tree.insert("hello", "tenant2")

        self.tree.insert("hey", "tenant2")  # NEWEST for t2

        # Verify initial sizes
        sizes_before = self.tree.get_used_size_per_tenant()
        self.assertEqual(sizes_before["tenant1"], 6)  # h(1) + i(1) + e(1) + llo(3) = 6
        self.assertEqual(
            sizes_before["tenant2"], 7
        )  # h(1) + i(1) + e(1) + llo(3) + y(1) = 7

        print("\nTree before eviction:")
        print(self.tree.pretty_print())

        # Evict - should remove "hi" from tenant2 as it's the oldest
        self.tree.evict_tenant_data(max_size)

        print("\nTree after eviction:")
        print(self.tree.pretty_print())

        # Verify sizes after eviction
        sizes_after = self.tree.get_used_size_per_tenant()
        self.assertEqual(sizes_after["tenant1"], 6)  # Should be unchanged
        self.assertEqual(sizes_after["tenant2"], 6)  # h(1) + e(1) + llo(3) + y(1) = 6

    def test_advanced_eviction(self):
        ...
        # Create 4 tenants
        # Each tenants keeps adding strings with shared prefixes to thousands usage
        # Set a strict limit for each tenant to only 100
        # At the end, check whether all of the tenant is under 100 after eviction

        max_size = {"tenant1": 100, "tenant2": 100, "tenant3": 100, "tenant4": 100}

        prefixes = ["aqwefcisdf", "iajsdfkmade", "kjnzxcvewqe", "iejksduqasd"]
        for i in range(100):
            for j, prefix in enumerate(prefixes):
                random_suffix = "".join(random.choices(string.ascii_letters, k=10))
                self.tree.insert(prefix + random_suffix, f"tenant{j+1}")

        sizes_before = self.tree.get_used_size_per_tenant()
        print(sizes_before)

        self.tree.evict_tenant_data(max_size)

        sizes_after = self.tree.get_used_size_per_tenant()
        print(sizes_after)
        # ensure size_after is below max_size
        for tenant, size in sizes_after.items():
            self.assertLessEqual(size, max_size[tenant])


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

	import random
	import string
	import time
	import unittest
	from typing import Dict, List, Tuple

	from tree import MultiTenantRadixTree


	class TestMultiTenantRadixTree(unittest.TestCase):
	def setUp(self):
	self.tree = MultiTenantRadixTree()

	def test_insert_exact_match(self):
	"""Test 1: Basic insert and exact match operations"""
	# Insert a single string for one tenant
	self.tree.insert("hello", "tenant1")
	matched, tenant = self.tree.prefix_match("hello")
	self.assertEqual(matched, "hello")
	self.assertEqual(tenant, "tenant1")

	# Insert same string for different tenant
	self.tree.insert("hello", "tenant2")
	matched, tenant = self.tree.prefix_match("hello")
	self.assertIn(tenant, ["tenant1", "tenant2"])

	# Insert different string for same tenant
	self.tree.insert("world", "tenant1")
	matched, tenant = self.tree.prefix_match("world")
	self.assertEqual(matched, "world")
	self.assertEqual(tenant, "tenant1")

	print(self.tree.pretty_print())

	def test_insert_partial_match(self):
	"""Test 2: Insert with partial matching scenarios"""
	# Test partial matches with common prefixes
	self.tree.insert("hello", "tenant1")
	print(self.tree.pretty_print())
	self.tree.insert("help", "tenant2")
	print(self.tree.pretty_print())

	# Match exact strings
	matched, tenant = self.tree.prefix_match("hello")
	self.assertEqual(matched, "hello")
	self.assertEqual(tenant, "tenant1")

	matched, tenant = self.tree.prefix_match("help")
	self.assertEqual(matched, "help")
	self.assertEqual(tenant, "tenant2")

	# Match partial string
	matched, tenant = self.tree.prefix_match("hel")
	self.assertEqual(matched, "hel")
	self.assertIn(tenant, ["tenant1", "tenant2"])

	# Match longer string
	matched, tenant = self.tree.prefix_match("hello_world")
	self.assertEqual(matched, "hello")
	self.assertEqual(tenant, "tenant1")

	def test_insert_edge_cases(self):
	"""Test 3: Edge cases for insert and match operations"""
	# Empty string
	self.tree.insert("", "tenant1")
	matched, tenant = self.tree.prefix_match("")
	self.assertEqual(matched, "")
	self.assertEqual(tenant, "tenant1")

	# Single character
	self.tree.insert("a", "tenant1")
	matched, tenant = self.tree.prefix_match("a")
	self.assertEqual(matched, "a")
	self.assertEqual(tenant, "tenant1")

	# Very long string
	long_str = "a" * 1000
	self.tree.insert(long_str, "tenant1")
	matched, tenant = self.tree.prefix_match(long_str)
	self.assertEqual(matched, long_str)
	self.assertEqual(tenant, "tenant1")

	# Unicode characters
	self.tree.insert("你好", "tenant1")
	matched, tenant = self.tree.prefix_match("你好")
	self.assertEqual(matched, "你好")
	self.assertEqual(tenant, "tenant1")

	def test_simple_eviction(self):
	"""Test 4: Simple eviction scenarios
	Tenant1: limit 10 chars
	Tenant2: limit 5 chars

	Should demonstrate:
	1. Basic eviction when size limit exceeded
	2. Proper eviction based on last access time
	3. Verification that shared nodes remain intact for other tenants
	"""
	# Set up size limits
	max_size = {"tenant1": 10, "tenant2": 5}

	# Insert strings for both tenants
	self.tree.insert("hello", "tenant1") # size 5
	self.tree.insert("hello", "tenant2") # size 5
	self.tree.insert("world", "tenant2") # size 5, total for tenant2 = 10

	# Verify initial sizes
	sizes_before = self.tree.get_used_size_per_tenant()
	self.assertEqual(sizes_before["tenant1"], 5) # "hello" = 5
	self.assertEqual(sizes_before["tenant2"], 10) # "hello" + "world" = 10

	# Evict - should remove "hello" from tenant2 as it's the oldest
	self.tree.evict_tenant_data(max_size)

	# Verify sizes after eviction
	sizes_after = self.tree.get_used_size_per_tenant()
	self.assertEqual(sizes_after["tenant1"], 5) # Should be unchanged
	self.assertEqual(sizes_after["tenant2"], 5) # Only "world" remains

	# Verify "world" remains for tenant2 (was accessed more recently)
	matched, tenant = self.tree.prefix_match("world")
	self.assertEqual(matched, "world")
	self.assertEqual(tenant, "tenant2")

	def test_medium_eviction(self):
	"""Test 5: Medium complexity eviction scenarios with shared prefixes
	Tenant1: limit 10 chars
	Tenant2: limit 7 chars (forces one string to be evicted)

	Tree structure after inserts:
	└── 'h' [t1, t2]
	├── 'i' [t1, t2] # Oldest for t2
	└── 'e' [t1, t2]
	├── 'llo' [t1, t2]
	└── 'y' [t2] # Newest for t2

	Size calculations:
	tenant1: "h"(1) + "i"(1) + "e"(1) + "llo"(3) = 6 chars
	tenant2: "h"(1) + "i"(1) + "e"(1) + "llo"(3) + "y"(1) = 7 chars

	After eviction (tenant2 exceeds limit by 1 char):
	"hi" should be removed from tenant2 as it's the oldest access
	"""
	max_size = {
	"tenant1": 10,
	"tenant2": 6,
	} # tenant2 will need to evict one string

	# Create a tree with overlapping prefixes
	self.tree.insert("hi", "tenant1")
	self.tree.insert("hi", "tenant2") # OLDEST for t2

	self.tree.insert("hello", "tenant1")
	self.tree.insert("hello", "tenant2")

	self.tree.insert("hey", "tenant2") # NEWEST for t2

	# Verify initial sizes
	sizes_before = self.tree.get_used_size_per_tenant()
	self.assertEqual(sizes_before["tenant1"], 6) # h(1) + i(1) + e(1) + llo(3) = 6
	self.assertEqual(
	sizes_before["tenant2"], 7
	) # h(1) + i(1) + e(1) + llo(3) + y(1) = 7

	print("\nTree before eviction:")
	print(self.tree.pretty_print())

	# Evict - should remove "hi" from tenant2 as it's the oldest
	self.tree.evict_tenant_data(max_size)

	print("\nTree after eviction:")
	print(self.tree.pretty_print())

	# Verify sizes after eviction
	sizes_after = self.tree.get_used_size_per_tenant()
	self.assertEqual(sizes_after["tenant1"], 6) # Should be unchanged
	self.assertEqual(sizes_after["tenant2"], 6) # h(1) + e(1) + llo(3) + y(1) = 6

	def test_advanced_eviction(self):
	...
	# Create 4 tenants
	# Each tenants keeps adding strings with shared prefixes to thousands usage
	# Set a strict limit for each tenant to only 100
	# At the end, check whether all of the tenant is under 100 after eviction

	max_size = {"tenant1": 100, "tenant2": 100, "tenant3": 100, "tenant4": 100}

	prefixes = ["aqwefcisdf", "iajsdfkmade", "kjnzxcvewqe", "iejksduqasd"]
	for i in range(100):
	for j, prefix in enumerate(prefixes):
	random_suffix = "".join(random.choices(string.ascii_letters, k=10))
	self.tree.insert(prefix + random_suffix, f"tenant{j+1}")

	sizes_before = self.tree.get_used_size_per_tenant()
	print(sizes_before)

	self.tree.evict_tenant_data(max_size)

	sizes_after = self.tree.get_used_size_per_tenant()
	print(sizes_after)
	# ensure size_after is below max_size
	for tenant, size in sizes_after.items():
	self.assertLessEqual(size, max_size[tenant])


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