diff --git a/ultravox/model/ultravox_config_test.py b/ultravox/model/ultravox_config_test.py index d00910d..4c40dd6 100644 --- a/ultravox/model/ultravox_config_test.py +++ b/ultravox/model/ultravox_config_test.py @@ -22,14 +22,13 @@ def test_can_load_release(model_id: str): config_from_dict = ultravox_config.UltravoxConfig(**orig_config.to_dict()) config_from_diff_dict = ultravox_config.UltravoxConfig(**orig_config.to_diff_dict()) keys_to_ignore = ("audio_latency_block_size",) + orig_values = { + **{k: None for k in keys_to_ignore}, + **orig_config.to_dict(), + } - assert ( - exclude_key(config_from_dict.to_dict(), keys_to_ignore) == orig_config.to_dict() - ) - assert ( - exclude_key(config_from_diff_dict.to_dict(), keys_to_ignore) - == orig_config.to_dict() - ) + assert config_from_dict.to_dict() == orig_values + assert config_from_diff_dict.to_dict() == orig_values assert config_from_dict.text_config.to_dict() == orig_config.text_config.to_dict() assert config_from_dict.audio_config.to_dict() == orig_config.audio_config.to_dict() @@ -38,13 +37,8 @@ def test_can_load_release(model_id: str): config_reloaded_diff = ultravox_config.UltravoxConfig( **config_from_dict.to_diff_dict() ) - assert ( - config_reloaded.to_dict() == {**{k: None for k in keys_to_ignore}, **orig_config.to_dict()} - ) - assert ( - exclude_key(config_reloaded_diff.to_dict(), keys_to_ignore) - == orig_config.to_dict() - ) + assert config_reloaded.to_dict() == orig_values + assert config_reloaded_diff.to_dict() == orig_values def test_no_config_when_id_present():