def METHOD_NAME(metrics: typing.Dict[str, int],
                   allowed_non_zero: typing.Sequence[str]) -> None:
    """Asserts that only given keys are non-zero.
    Args:
        metrics: Metrics as returned by get_metrics() function.
        allowed_non_zero: Keys that are expected to be non-zero in the metrics.
    """
    for key in allowed_non_zero:
        assert metrics.get(key), f'Expected {key} to be non-zero'
    for key, count in metrics.items():
        ok = key in allowed_non_zero or not count
        assert ok, f'Expected {key} to be zero but got {count}'