def test_values_as_expected(self):
    """Test whether processed dictionaries match expected values."""

    metric_dict = {
        'cat1': {
            'key': 2.0,
            'denominator': 1.0
        },
        'cat2': {
            'key': 2.0,
            'denominator': 2.0
        },
    }

    processed_metrics = metric_utils.process_metrics(metric_dict)  # pytype: disable=wrong-arg-types  # jax-ndarray
    expected_result = {
        'cat1_key': 2.0,
        'cat1_denom': 1.0,
        'cat2_key': 1.0,
        'cat2_denom': 2.0,
    }
    self.assertEqual(processed_metrics, expected_result)

    metric_dict = {
        'cat1': {
            'key': 2.0,
            'denominator': 1.0
        },
        'cat2': {
            'key': 2.0,
            'denominator': 2.0
        },
    }

    processed_metrics = metric_utils.process_metrics(metric_dict, prefix='pref')  # pytype: disable=wrong-arg-types  # jax-ndarray
    expected_result = {
        'pref/cat1_key': 2.0,
        'pref/cat1_denom': 1.0,
        'pref/cat2_key': 1.0,
        'pref/cat2_denom': 2.0,
    }
    self.assertEqual(processed_metrics, expected_result)
Test whether processed dictionaries match expected values.
language/mentionmemory/utils/metric_utils_test.py