Coverage for NeuralTSNE/NeuralTSNE/Utils/Loaders/LabelLoaders/tests/test_label_loaders.py: 100%

1import io 

2 

3import pytest 

4import torch 

5 

6from NeuralTSNE.Utils.Loaders.LabelLoaders import load_labels 

7 

8 

9@pytest.mark.parametrize( 

10 "labels", ["1\n2\n3\n", "3\n2\n1\n6\n", "1\n3\n2\n", "2\n1\n3\n4\n", None] 

11) 

12def test_load_labels(labels: str | None): 

13 if labels is None: 

14 assert load_labels(labels) is None 

15 else: 

16 expected = torch.tensor([float(label) for label in labels.splitlines()]) 

17 labels_file = io.StringIO(labels) 

18 assert torch.allclose(load_labels(labels_file), expected) 

19 assert labels_file.closed