diff --git a/amqp/config/amq_configuration.py b/amqp/config/amq_configuration.py index 1ed3113..fe9faf2 100644 --- a/amqp/config/amq_configuration.py +++ b/amqp/config/amq_configuration.py @@ -88,8 +88,21 @@ class AMQAdapter: "CleverMicro-AMQ", self.PREFIX + "consul-initial-retry-delay", fallback=0.2 ) ) - self.consul_counter_key = config.get( - "CleverMicro-AMQ", self.PREFIX + "consul-counter-key", fallback="service/ids/counter" + self.consul_counter_key = ( + os.getenv( + "CONSUL_COUNTER_KEY", + config.get( + "CleverMicro-AMQ", + self.PREFIX + "consul-counter-key", + fallback="service/ids/counter", + ), + ) + if os.getenv("CONSUL_COUNTER_KEY") is not None + else config.get( + "CleverMicro-AMQ", + self.PREFIX + "consul-counter-key", + fallback="service/ids/counter", + ) ) def adapter_prefix(self) -> str: diff --git a/tests/adapter/test_consul_global_id_generator.py b/tests/adapter/test_consul_global_id_generator.py new file mode 100644 index 0000000..acb5087 --- /dev/null +++ b/tests/adapter/test_consul_global_id_generator.py @@ -0,0 +1,203 @@ +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()