| # 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, | |
| # software 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() |