diff --git a/concordia/utils/async_measurements_test.py b/concordia/utils/async_measurements_test.py new file mode 100644 index 000000000..183e480d6 --- /dev/null +++ b/concordia/utils/async_measurements_test.py @@ -0,0 +1,90 @@ +# Copyright 2025 DeepMind Technologies Limited. +# +# Licensed 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 +# +# https://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. + +"""Tests for ReactiveMeasurements.""" + +from absl.testing import absltest +from concordia.utils import async_measurements + + +class ReactiveMeasurementsTest(absltest.TestCase): + + def test_publish_datum_still_stores_in_channel(self): + m = async_measurements.ReactiveMeasurements() + m.publish_datum('a', 1) + self.assertEqual(m.get_channel('a'), [1]) + + def test_subscribe_receives_published_data(self): + m = async_measurements.ReactiveMeasurements() + received = [] + subscription = m.subscribe(lambda pair: received.append(pair)) + m.publish_datum('a', 1) + m.publish_datum('b', 2) + subscription.dispose() + self.assertEqual(received, [('a', 1), ('b', 2)]) + + def test_capture_collects_only_data_with_matching_key(self): + m = async_measurements.ReactiveMeasurements() + with m.capture('entity_1') as captured: + m.publish_datum('a', 'mine', capture_key='entity_1') + m.publish_datum('b', 'not_mine', capture_key='entity_2') + m.publish_datum('c', 'uncaptured') + self.assertEqual(captured, {'a': 'mine'}) + + def test_capture_only_keeps_last_datum_per_channel(self): + m = async_measurements.ReactiveMeasurements() + with m.capture('entity_1') as captured: + m.publish_datum('a', 'first', capture_key='entity_1') + m.publish_datum('a', 'second', capture_key='entity_1') + self.assertEqual(captured, {'a': 'second'}) + + def test_capture_stops_collecting_after_context_exits(self): + m = async_measurements.ReactiveMeasurements() + with m.capture('entity_1') as captured: + m.publish_datum('a', 'inside', capture_key='entity_1') + m.publish_datum('a', 'outside', capture_key='entity_1') + self.assertEqual(captured, {'a': 'inside'}) + + def test_nested_captures_with_different_keys_do_not_cross_contaminate(self): + m = async_measurements.ReactiveMeasurements() + with m.capture('entity_1') as captured_1: + with m.capture('entity_2') as captured_2: + m.publish_datum('a', 'for_1', capture_key='entity_1') + m.publish_datum('a', 'for_2', capture_key='entity_2') + self.assertEqual(captured_1, {'a': 'for_1'}) + self.assertEqual(captured_2, {'a': 'for_2'}) + + def test_close_still_works_on_reactive_subclass(self): + # Regression coverage: close()/close_channel() must not deadlock on the + # ReactiveMeasurements subclass either. + m = async_measurements.ReactiveMeasurements() + m.publish_datum('a', 1) + m.publish_datum('b', 2) + m.close() + self.assertEqual(m.available_channels(), set()) + + def test_dispose_completes_subject(self): + m = async_measurements.ReactiveMeasurements() + received = [] + completed = [] + m.subscribe(lambda pair: received.append(pair)) + m._subject.subscribe(on_completed=lambda: completed.append(True)) # pylint: disable=protected-access + m.publish_datum('a', 1) + m.dispose() + self.assertEqual(received, [('a', 1)]) + self.assertTrue(completed) + + +if __name__ == '__main__': + absltest.main() diff --git a/concordia/utils/measurements.py b/concordia/utils/measurements.py index e02ec4ea2..52195c1d4 100644 --- a/concordia/utils/measurements.py +++ b/concordia/utils/measurements.py @@ -97,15 +97,13 @@ def close_channel(self, channel: str) -> None: """Closes the channel for the given name. Args: - channel: The channel to close. If the channel doesn't exist yet, it will - be created. + channel: The channel to close. If the channel doesn't exist, this is a + no-op. """ with self._channels_lock: - del self._channels[channel] + self._channels.pop(channel, None) def close(self) -> None: """Closes all channels.""" with self._channels_lock: - for channel in self._channels: - self.close_channel(channel) self._channels.clear() diff --git a/concordia/utils/measurements_test.py b/concordia/utils/measurements_test.py new file mode 100644 index 000000000..5538146d7 --- /dev/null +++ b/concordia/utils/measurements_test.py @@ -0,0 +1,117 @@ +# Copyright 2024 DeepMind Technologies Limited. +# +# Licensed 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 +# +# https://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. + +"""Tests for the Measurements registry.""" + +import threading + +from absl.testing import absltest +from concordia.utils import measurements as measurements_lib + + +class MeasurementsTest(absltest.TestCase): + + def test_publish_datum_creates_channel(self): + m = measurements_lib.Measurements() + m.publish_datum('channel_a', 1) + self.assertEqual(m.available_channels(), {'channel_a'}) + + def test_publish_datum_appends_in_order(self): + m = measurements_lib.Measurements() + m.publish_datum('channel_a', 1) + m.publish_datum('channel_a', 2) + m.publish_datum('channel_a', 3) + self.assertEqual(m.get_channel('channel_a'), [1, 2, 3]) + + def test_get_channel_creates_empty_channel_if_missing(self): + m = measurements_lib.Measurements() + self.assertEqual(m.get_channel('missing'), []) + self.assertEqual(m.available_channels(), {'missing'}) + + def test_get_last_datum_returns_most_recent(self): + m = measurements_lib.Measurements() + m.publish_datum('channel_a', 'first') + m.publish_datum('channel_a', 'second') + self.assertEqual(m.get_last_datum('channel_a'), 'second') + + def test_get_last_datum_returns_none_for_empty_channel(self): + m = measurements_lib.Measurements() + self.assertIsNone(m.get_last_datum('missing')) + + def test_get_all_channels_returns_all_data(self): + m = measurements_lib.Measurements() + m.publish_datum('a', 1) + m.publish_datum('b', 2) + self.assertEqual(m.get_all_channels(), {'a': [1], 'b': [2]}) + + def test_get_all_channels_returns_a_copy(self): + m = measurements_lib.Measurements() + m.publish_datum('a', 1) + snapshot = m.get_all_channels() + m.publish_datum('a', 2) + self.assertEqual(snapshot, {'a': [1]}) + self.assertEqual(m.get_channel('a'), [1, 2]) + + def test_close_channel_removes_channel(self): + m = measurements_lib.Measurements() + m.publish_datum('a', 1) + m.close_channel('a') + self.assertEqual(m.available_channels(), set()) + + def test_close_channel_on_missing_channel_is_a_no_op(self): + m = measurements_lib.Measurements() + # Should not raise, per the method's documented behavior. + m.close_channel('never_published') + self.assertEqual(m.available_channels(), set()) + + def test_close_removes_all_channels(self): + m = measurements_lib.Measurements() + m.publish_datum('a', 1) + m.publish_datum('b', 2) + m.publish_datum('c', 3) + m.close() + self.assertEqual(m.available_channels(), set()) + + def test_close_does_not_deadlock_with_multiple_channels(self): + # Regression test: close() used to hold _channels_lock while calling + # close_channel(), which itself acquired the same non-reentrant lock, + # deadlocking whenever there was at least one channel. + m = measurements_lib.Measurements() + m.publish_datum('a', 1) + m.publish_datum('b', 2) + + completed = threading.Event() + + def run_close(): + m.close() + completed.set() + + thread = threading.Thread(target=run_close) + thread.start() + thread.join(timeout=5) + self.assertTrue(completed.is_set(), 'Measurements.close() deadlocked.') + + def test_close_on_empty_measurements_does_not_raise(self): + m = measurements_lib.Measurements() + m.close() + self.assertEqual(m.available_channels(), set()) + + def test_get_channel_or_create_requires_lock(self): + m = measurements_lib.Measurements() + with self.assertRaises(RuntimeError): + m._get_channel_or_create('a') # pylint: disable=protected-access + + +if __name__ == '__main__': + absltest.main()