| ################################################################################ |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| ################################################################################ |
| import os |
| import tempfile |
| import time |
| import unittest |
| from datetime import datetime |
| |
| from pypaimon.consumer.consumer import Consumer |
| from pypaimon.consumer.consumer_manager import ConsumerManager |
| from pypaimon.filesystem.local_file_io import LocalFileIO |
| |
| |
| class ConsumerManagerTest(unittest.TestCase): |
| """Test cases for ConsumerManager.""" |
| |
| def setUp(self): |
| """Set up test environment.""" |
| self.temp_dir = tempfile.mkdtemp() |
| self.file_io = LocalFileIO(self.temp_dir) |
| self.manager = ConsumerManager(self.file_io, self.temp_dir) |
| self.consumer_manager_branch = ConsumerManager(self.file_io, self.temp_dir, "branch1") |
| |
| def tearDown(self): |
| """Clean up test environment.""" |
| import shutil |
| if os.path.exists(self.temp_dir): |
| shutil.rmtree(self.temp_dir) |
| |
| def test_retry(self): |
| """Test retry mechanism for corrupted consumer file.""" |
| # Create corrupted consumer file |
| consumer_dir = os.path.join(self.temp_dir, "consumer") |
| os.makedirs(consumer_dir, exist_ok=True) |
| consumer_file = os.path.join(consumer_dir, "consumer-id1") |
| with open(consumer_file, 'w') as f: |
| f.write("invalid json content") |
| |
| # Should raise RuntimeError after retries |
| with self.assertRaises(RuntimeError) as context: |
| self.manager.consumer("id1") |
| self.assertIn("Retry fail after 10 times", str(context.exception)) |
| |
| def test_basic_operations(self): |
| """Test basic consumer operations.""" |
| # Test non-existent consumer |
| consumer = self.manager.consumer("id1") |
| self.assertIsNone(consumer) |
| |
| # Test min next snapshot when no consumers |
| min_snapshot = self.manager.min_next_snapshot() |
| self.assertIsNone(min_snapshot) |
| |
| # Reset consumer |
| self.manager.reset_consumer("id1", Consumer(5)) |
| consumer = self.manager.consumer("id1") |
| self.assertIsNotNone(consumer) |
| self.assertEqual(consumer.next_snapshot, 5) |
| |
| # Reset another consumer |
| self.manager.reset_consumer("id2", Consumer(8)) |
| consumer = self.manager.consumer("id2") |
| self.assertIsNotNone(consumer) |
| self.assertEqual(consumer.next_snapshot, 8) |
| |
| # Test min next snapshot |
| min_snapshot = self.manager.min_next_snapshot() |
| self.assertEqual(min_snapshot, 5) |
| |
| def test_branch_operations(self): |
| """Test consumer operations on different branches.""" |
| # Test non-existent consumer on branch |
| consumer = self.consumer_manager_branch.consumer("id1") |
| self.assertIsNone(consumer) |
| |
| # Test min next snapshot when no consumers on branch |
| min_snapshot = self.consumer_manager_branch.min_next_snapshot() |
| self.assertIsNone(min_snapshot) |
| |
| # Reset consumer on branch |
| self.consumer_manager_branch.reset_consumer("id1", Consumer(5)) |
| consumer = self.consumer_manager_branch.consumer("id1") |
| self.assertIsNotNone(consumer) |
| self.assertEqual(consumer.next_snapshot, 5) |
| |
| # Reset another consumer on branch |
| self.consumer_manager_branch.reset_consumer("id2", Consumer(8)) |
| consumer = self.consumer_manager_branch.consumer("id2") |
| self.assertIsNotNone(consumer) |
| self.assertEqual(consumer.next_snapshot, 8) |
| |
| # Test min next snapshot on branch |
| min_snapshot = self.consumer_manager_branch.min_next_snapshot() |
| self.assertEqual(min_snapshot, 5) |
| |
| def test_expire(self): |
| """Test consumer expiration.""" |
| # Create consumers with different timestamps |
| self.manager.reset_consumer("id1", Consumer(1)) |
| time.sleep(1) |
| expire_datetime = datetime.now() |
| time.sleep(1) |
| self.manager.reset_consumer("id2", Consumer(2)) |
| |
| # Check expire |
| self.manager.expire(expire_datetime) |
| consumer1 = self.manager.consumer("id1") |
| self.assertIsNone(consumer1) |
| consumer2 = self.manager.consumer("id2") |
| self.assertIsNotNone(consumer2) |
| self.assertEqual(consumer2.next_snapshot, 2) |
| |
| # Check last modification |
| expire_datetime = datetime.now() |
| time.sleep(1) |
| self.manager.reset_consumer("id2", Consumer(3)) |
| self.manager.expire(expire_datetime) |
| consumer2 = self.manager.consumer("id2") |
| self.assertIsNotNone(consumer2) |
| self.assertEqual(consumer2.next_snapshot, 3) |
| |
| def test_expire_branch(self): |
| """Test consumer expiration on branch.""" |
| # Create consumers on branch with different timestamps |
| self.consumer_manager_branch.reset_consumer("id3", Consumer(1)) |
| time.sleep(1) |
| expire_datetime = datetime.now() |
| time.sleep(1) |
| self.consumer_manager_branch.reset_consumer("id4", Consumer(2)) |
| |
| # Check expire on branch |
| self.consumer_manager_branch.expire(expire_datetime) |
| consumer3 = self.consumer_manager_branch.consumer("id3") |
| self.assertIsNone(consumer3) |
| consumer4 = self.consumer_manager_branch.consumer("id4") |
| self.assertIsNotNone(consumer4) |
| self.assertEqual(consumer4.next_snapshot, 2) |
| |
| # Check last modification on branch |
| expire_datetime = datetime.now() |
| time.sleep(1) |
| self.consumer_manager_branch.reset_consumer("id4", Consumer(3)) |
| self.consumer_manager_branch.expire(expire_datetime) |
| consumer4 = self.consumer_manager_branch.consumer("id4") |
| self.assertIsNotNone(consumer4) |
| self.assertEqual(consumer4.next_snapshot, 3) |
| |
| def test_read_consumer(self): |
| """Test reading consumer from different branches.""" |
| # Create consumer on main branch |
| self.manager.reset_consumer("id1", Consumer(5)) |
| consumer = self.manager.consumer("id1") |
| self.assertIsNotNone(consumer) |
| self.assertEqual(consumer.next_snapshot, 5) |
| |
| # Create consumer on branch |
| self.consumer_manager_branch.reset_consumer("id2", Consumer(5)) |
| consumer = self.consumer_manager_branch.consumer("id2") |
| self.assertIsNotNone(consumer) |
| self.assertEqual(consumer.next_snapshot, 5) |
| |
| # Verify id2 doesn't exist on main branch |
| consumer = self.manager.consumer("id2") |
| self.assertIsNone(consumer) |
| |
| def test_list_all_ids(self): |
| """Test listing all consumer IDs.""" |
| # Initially empty |
| ids = self.manager.list_all_ids() |
| self.assertEqual(len(ids), 0) |
| |
| # Add consumers |
| self.manager.reset_consumer("id1", Consumer(1)) |
| self.manager.reset_consumer("id2", Consumer(2)) |
| self.manager.reset_consumer("id3", Consumer(3)) |
| |
| # List IDs |
| ids = self.manager.list_all_ids() |
| self.assertEqual(len(ids), 3) |
| self.assertIn("id1", ids) |
| self.assertIn("id2", ids) |
| self.assertIn("id3", ids) |
| |
| def test_consumers(self): |
| """Test getting all consumers.""" |
| # Initially empty |
| consumers = self.manager.consumers() |
| self.assertEqual(len(consumers), 0) |
| |
| # Add consumers |
| self.manager.reset_consumer("id1", Consumer(1)) |
| self.manager.reset_consumer("id2", Consumer(2)) |
| self.manager.reset_consumer("id3", Consumer(3)) |
| |
| # Get all consumers |
| consumers = self.manager.consumers() |
| self.assertEqual(len(consumers), 3) |
| self.assertEqual(consumers["id1"], 1) |
| self.assertEqual(consumers["id2"], 2) |
| self.assertEqual(consumers["id3"], 3) |
| |
| def test_delete_consumer(self): |
| """Test deleting consumer.""" |
| # Create consumer |
| self.manager.reset_consumer("id1", Consumer(5)) |
| consumer = self.manager.consumer("id1") |
| self.assertIsNotNone(consumer) |
| |
| # Delete consumer |
| self.manager.delete_consumer("id1") |
| consumer = self.manager.consumer("id1") |
| self.assertIsNone(consumer) |
| |
| def test_clear_consumers(self): |
| """Test clearing consumers with patterns.""" |
| # Add multiple consumers |
| self.manager.reset_consumer("test-id1", Consumer(1)) |
| self.manager.reset_consumer("test-id2", Consumer(2)) |
| self.manager.reset_consumer("prod-id1", Consumer(3)) |
| self.manager.reset_consumer("prod-id2", Consumer(4)) |
| |
| # Clear test consumers |
| self.manager.clear_consumers("test-.*") |
| |
| # Verify test consumers are deleted |
| self.assertIsNone(self.manager.consumer("test-id1")) |
| self.assertIsNone(self.manager.consumer("test-id2")) |
| |
| # Verify prod consumers remain |
| self.assertIsNotNone(self.manager.consumer("prod-id1")) |
| self.assertIsNotNone(self.manager.consumer("prod-id2")) |
| |
| def test_clear_consumers_with_exclusion(self): |
| """Test clearing consumers with inclusion and exclusion patterns.""" |
| # Add multiple consumers |
| self.manager.reset_consumer("test-id1", Consumer(1)) |
| self.manager.reset_consumer("test-id2", Consumer(2)) |
| self.manager.reset_consumer("test-backup", Consumer(3)) |
| |
| # Clear test consumers but exclude backup |
| self.manager.clear_consumers("test-.*", "test-backup") |
| |
| # Verify test-id1 and test-id2 are deleted |
| self.assertIsNone(self.manager.consumer("test-id1")) |
| self.assertIsNone(self.manager.consumer("test-id2")) |
| |
| # Verify test-backup remains |
| self.assertIsNotNone(self.manager.consumer("test-backup")) |
| |
| def test_with_branch(self): |
| """Test with_branch method.""" |
| # Create consumer on main branch |
| self.manager.reset_consumer("main_consumer", Consumer(10)) |
| main_consumer = self.manager.consumer("main_consumer") |
| self.assertIsNotNone(main_consumer) |
| self.assertEqual(main_consumer.next_snapshot, 10) |
| |
| # Create consumer manager for a different branch using with_branch |
| branch_manager = self.manager.with_branch("feature_branch") |
| self.assertEqual(branch_manager._branch, "feature_branch") |
| |
| # Verify main branch consumer doesn't exist on new branch |
| branch_consumer = branch_manager.consumer("main_consumer") |
| self.assertIsNone(branch_consumer) |
| |
| # Create consumer on new branch |
| branch_manager.reset_consumer("branch_consumer", Consumer(20)) |
| branch_consumer = branch_manager.consumer("branch_consumer") |
| self.assertIsNotNone(branch_consumer) |
| self.assertEqual(branch_consumer.next_snapshot, 20) |
| |
| # Verify original manager still operates on main branch |
| self.assertEqual(self.manager._branch, "main") |
| main_consumer = self.manager.consumer("main_consumer") |
| self.assertIsNotNone(main_consumer) |
| self.assertEqual(main_consumer.next_snapshot, 10) |
| |
| # Verify branch consumer doesn't exist on main branch |
| main_branch_consumer = self.manager.consumer("branch_consumer") |
| self.assertIsNone(main_branch_consumer) |
| |
| def test_with_branch_main(self): |
| """Test with_branch with None returns main branch.""" |
| branch_manager = self.manager.with_branch(None) |
| self.assertEqual(branch_manager._branch, "main") |
| |
| def test_with_branch_empty(self): |
| """Test with_branch with empty string returns main branch.""" |
| branch_manager = self.manager.with_branch("") |
| self.assertEqual(branch_manager._branch, "main") |
| |
| def test_with_branch_explicit_main(self): |
| """Test with_branch with explicit 'main' branch.""" |
| branch_manager = self.manager.with_branch("main") |
| self.assertEqual(branch_manager._branch, "main") |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
