import unittest from unittest.mock import MagicMock, call, patch from amqp.adapter.consul_global_id_generator import ( generate_fallback_id, get_config_values, get_unique_instance_id, ) from amqp.config.amq_configuration import AMQConfiguration class TestConsulGlobalIdGenerator(unittest.TestCase): def setUp(self): self.mock_config = MagicMock() self.mock_config.consul_host = "test-consul" self.mock_config.consul_port = 8500 self.mock_config.consul_counter_key = "test/counter/key" self.mock_config.consul_max_retries = 3 self.mock_config.consul_initial_retry_delay = 0.1 def test_get_config_values_from_config(self): """Test retrieving configuration values from AMQAdapter config.""" host, port, key, retries, delay = get_config_values(self.mock_config) self.assertEqual(host, "test-consul") self.assertEqual(port, 8500) self.assertEqual(key, "test/counter/key") self.assertEqual(retries, 3) self.assertEqual(delay, 0.1) @patch.dict( "os.environ", { "CONSUL_HOST": "env-consul", "CONSUL_PORT": "8501", "CONSUL_COUNTER_KEY": "env/counter/key", "CONSUL_MAX_RETRIES": "4", "CONSUL_RETRY_DELAY_SECONDS": "0.2", }, ) def test_get_config_values_from_env(self): """Test retrieving configuration values from environment variables when config fails.""" self.mock_config = MagicMock(side_effect=Exception("Config error")) host, port, key, retries, delay = get_config_values(AMQConfiguration("").amq_adapter) self.assertEqual(host, "env-consul") self.assertEqual(port, 8501) self.assertEqual(key, "env/counter/key") self.assertEqual(retries, 4) self.assertEqual(delay, 0.2) @patch("time.time", return_value=1000.0) @patch("socket.gethostname", return_value="test-host") @patch("random.randint", return_value=1234) def test_generate_fallback_id(self, mock_randint, mock_hostname, mock_time): """Test fallback ID generation with controlled inputs.""" # Calculate expected value based on the mocked values timestamp = int(1000.0 * 1000) hostname_hash = hash("test-host") % 10000 random_part = 1234 expected_id = (timestamp * 100000000) + (hostname_hash * 10000) + random_part result = generate_fallback_id() self.assertEqual(result, expected_id) @patch("consul_kv.Connection") def test_get_unique_instance_id_success(self, mock_connection): """Test successful ID retrieval from Consul.""" mock_consul = MagicMock() mock_connection.return_value = mock_consul # Mock the get and put methods mock_data = {"Value": b"42", "ModifyIndex": 123} mock_consul.get.return_value = (0, mock_data) mock_consul.put.return_value = True result = get_unique_instance_id(self.mock_config) # Should return the incremented value self.assertEqual(result, 43) mock_connection.assert_called_once_with(endpoint="test-consul:8500", timeout=5) mock_consul.get.assert_called_once_with("test/counter/key") mock_consul.put.assert_called_once_with("test/counter/key", "43", cas=123) @patch("consul_kv.Connection") def test_get_unique_instance_id_initialize(self, mock_connection): """Test initializing the counter when it doesn't exist.""" mock_consul = MagicMock() mock_connection.return_value = mock_consul # Return None to simulate counter not existing mock_consul.get.return_value = (0, None) mock_consul.put.return_value = True result = get_unique_instance_id(self.mock_config) # Should return 1 for the first ID self.assertEqual(result, 1) mock_consul.put.assert_called_once_with("test/counter/key", "1", cas=1) @patch("consul_kv.Connection") @patch("amqp.adapter.consul_global_id_generator.generate_fallback_id", return_value=9999) def test_get_unique_instance_id_initialize_failure(self, mock_fallback, mock_connection): """Test fallback when initializing the counter fails.""" mock_consul = MagicMock() mock_connection.return_value = mock_consul # Return None to simulate counter not existing mock_consul.get.return_value = (0, None) mock_consul.put.return_value = False # Initialization fails result = get_unique_instance_id(self.mock_config) # Should return fallback ID self.assertEqual(result, 9999) mock_fallback.assert_called_once() @patch("consul_kv.Connection") @patch("time.sleep") def test_get_unique_instance_id_retry_success(self, mock_sleep, mock_connection): """Test successful retry after CAS failure.""" mock_consul = MagicMock() mock_connection.return_value = mock_consul # First attempt fails, second succeeds mock_data1 = {"Value": b"42", "ModifyIndex": 123} mock_data2 = {"Value": b"43", "ModifyIndex": 124} # First put fails, second succeeds mock_consul.get.side_effect = [(0, mock_data1), (0, mock_data2)] mock_consul.put.side_effect = [False, True] result = get_unique_instance_id(self.mock_config) # Should return the incremented value from the second attempt self.assertEqual(result, 44) self.assertEqual(mock_consul.get.call_count, 2) self.assertEqual(mock_consul.put.call_count, 2) mock_sleep.assert_called_once_with(0.1) # First retry delay @patch("consul_kv.Connection") @patch("time.sleep") @patch("amqp.adapter.consul_global_id_generator.generate_fallback_id", return_value=9999) def test_get_unique_instance_id_max_retries(self, mock_fallback, mock_sleep, mock_connection): """Test fallback after max retries.""" mock_consul = MagicMock() mock_connection.return_value = mock_consul # All attempts fail mock_data = {"Value": b"42", "ModifyIndex": 123} mock_consul.get.return_value = (0, mock_data) mock_consul.put.return_value = False result = get_unique_instance_id(self.mock_config) # Should return fallback ID self.assertEqual(result, 9999) self.assertEqual(mock_consul.put.call_count, 3) # 3 retries as configured mock_fallback.assert_called_once() # Check exponential backoff expected_calls = [ call(0.1), # First retry call(0.2), # Second retry (doubled) ] mock_sleep.assert_has_calls(expected_calls) @patch("consul_kv.Connection") @patch("amqp.adapter.consul_global_id_generator.generate_fallback_id", return_value=9999) def test_get_unique_instance_id_consul_exception(self, mock_fallback, mock_connection): """Test fallback when Consul raises an exception.""" mock_connection.side_effect = Exception("Consul connection error") result = get_unique_instance_id(self.mock_config) # Should return fallback ID self.assertEqual(result, 9999) mock_fallback.assert_called_once() @patch("consul_kv.Connection") @patch("amqp.adapter.consul_global_id_generator.generate_fallback_id", return_value=9999) def test_get_unique_instance_id_counter_disappeared(self, mock_fallback, mock_connection): """Test fallback when counter disappears during retry.""" mock_consul = MagicMock() mock_connection.return_value = mock_consul # First get succeeds, second returns None (counter disappeared) mock_data = {"Value": b"42", "ModifyIndex": 123} mock_consul.get.side_effect = [(0, mock_data), (0, None)] mock_consul.put.return_value = False result = get_unique_instance_id(self.mock_config) # Should return fallback ID self.assertEqual(result, 9999) mock_fallback.assert_called_once() if __name__ == "__main__": unittest.main()