Repository: Lightning-AI/pytorch-lightning
File: tests/utilities/parsing.py (Python)

Code before:
    return model1, model2, model3, model4, model5


def test_lightning_hasattr(tmpdir):
    """ Test that the lightning_hasattr works in all cases"""
    model1, model2, model3, model4, model5 = _get_test_cases()
    assert lightning_hasattr(model1, 'learning_rate'), \
        'lightning_hasattr failed to find namespace variable'
    assert lightning_hasattr(model2, 'learning_rate'), \
        'lightning_hasattr failed to find hparams namespace variable'
    assert lightning_hasattr(model3, 'learning_rate'), \
        'lightning_hasattr failed to find hparams dict variable'
    assert not lightning_hasattr(model4, 'learning_rate'), \
        'lightning_hasattr found variable when it should not'
    assert lightning_hasattr(model5, 'batch_size'), \
        'lightning_hasattr failed to find batch_size in datamodule'


def test_lightning_getattr(tmpdir):
    """ Test that the lightning_getattr works in all cases"""
    models = _get_test_cases()
    for i, m in enumerate(models[:3]):
        value = lightning_getattr(m, 'learning_rate')
        assert value == i, 'attribute not correctly extracted'

    model5 = models[-1]
    assert lightning_getattr(model5, 'batch_size') == 8, \
        'batch_size not correctly extracted'


def test_lightning_setattr(tmpdir):
    """ Test that the lightning_setattr works in all cases"""
    models = _get_test_cases()
    for m in models[:3]:
        lightning_setattr(m, 'learning_rate', 10)
        assert lightning_getattr(m, 'learning_rate') == 10, \
            'attribute not correctly set'

    model5 = models[-1]
    lightning_setattr(model5, 'batch_size', 128)
    assert lightning_getattr(model5, 'batch_size') == 128, \
        'batch_size not correctly set'


Review comment (suggestion):
```suggestion
    model5 = models[4]
```
or it will potentially break if a new model is added.

Code after:
    return model1, model2, model3, model4, model5


def test_lightning_hasattr(tmpdir):
    """ Test that the lightning_hasattr works in all cases"""
    model1, model2, model3, model4, model5 = _get_test_cases()
    assert lightning_hasattr(model1, 'learning_rate'), \
        'lightning_hasattr failed to find namespace variable'
    assert lightning_hasattr(model2, 'learning_rate'), \
        'lightning_hasattr failed to find hparams namespace variable'
    assert lightning_hasattr(model3, 'learning_rate'), \
        'lightning_hasattr failed to find hparams dict variable'
    assert not lightning_hasattr(model4, 'learning_rate'), \
        'lightning_hasattr found variable when it should not'
    assert lightning_hasattr(model5, 'batch_size'), \
        'lightning_hasattr failed to find batch_size in datamodule'


def test_lightning_getattr(tmpdir):
    """ Test that the lightning_getattr works in all cases"""
    models = _get_test_cases()
    for i, m in enumerate(models[:3]):
        value = lightning_getattr(m, 'learning_rate')
        assert value == i, 'attribute not correctly extracted'

    model5 = models[4]
    assert lightning_getattr(model5, 'batch_size') == 8, \
        'batch_size not correctly extracted'


def test_lightning_setattr(tmpdir):
    """ Test that the lightning_setattr works in all cases"""
    models = _get_test_cases()
    for m in models[:3]:
        lightning_setattr(m, 'learning_rate', 10)
        assert lightning_getattr(m, 'learning_rate') == 10, \
            'attribute not correctly set'

    model5 = models[4]
    lightning_setattr(model5, 'batch_size', 128)
    assert lightning_getattr(model5, 'batch_size') == 128, \
        'batch_size not correctly set'