test: Improve test ids (#460)

* Move conftest.py from tests.unit to tests
* Separate DEVICE creation from conftest.py to device.py
* Add pytest_make_parametrize_id

Author: ValerianRey

CONTRIBUTING.md

@@ -130,18 +130,19 @@ from utils.tensors import ones_
 a = ones_(3, 4)
 ```
 
-This will automatically call `torch.ones` with `device=unit.conftest.DEVICE`.
+This will automatically call `torch.ones` with `device=DEVICE`.
 If the function you need does not exist yet as a partial function in `tensors.py`, add it.
 Lastly, when you create a model or a random generator, you have to move them manually to the right
-device (the `DEVICE` defined in `unit.conftest`):
+device (the `DEVICE` defined in `device.py`).
 ```python
 import torch
 from torch.nn import Linear
-from unit.conftest import DEVICE
+from device import DEVICE
 
 model = Linear(3, 4).to(device=DEVICE)
 rng = torch.Generator(device=DEVICE)
 ```
+You may also use a `ModuleFactory` to make the modules on `DEVICE` automatically.
 
 ### Coding
 

tests/unit/conftest.py renamed to tests/conftest.py

@@ -1,21 +1,13 @@
-import os
 import random as rand
+from contextlib import nullcontext
 
 import torch
-from pytest import fixture, mark
+from device import DEVICE
+from pytest import RaisesExc, fixture, mark
+from torch import Tensor
+from utils.architectures import ModuleFactory
 
-try:
-    _device_str = os.environ["PYTEST_TORCH_DEVICE"]
-except KeyError:
-    _device_str = "cpu"  # Default to cpu if environment variable not set
-
-if _device_str != "cuda:0" and _device_str != "cpu":
-    raise ValueError(f"Invalid value of environment variable PYTEST_TORCH_DEVICE: {_device_str}")
-
-if _device_str == "cuda:0" and not torch.cuda.is_available():
-    raise ValueError('Requested device "cuda:0" but cuda is not available.')
-
-DEVICE = torch.device(_device_str)
+from torchjd.aggregation import Aggregator, Weighting
 
 
 @fixture(autouse=True)
@@ -48,3 +40,24 @@ def pytest_collection_modifyitems(config, items):
             item.add_marker(skip_slow)
         if "xfail_if_cuda" in item.keywords and str(DEVICE).startswith("cuda"):
             item.add_marker(xfail_cuda)
+
+
+def pytest_make_parametrize_id(config, val, argname):
+    MAX_SIZE = 40
+    optional_string = None  # Returning None means using pytest's way of making the string
+
+    if isinstance(val, (Aggregator, ModuleFactory, Weighting)):
+        optional_string = str(val)
+    elif isinstance(val, Tensor):
+        optional_string = "T" + str(list(val.shape))  # T to indicate that it's a tensor
+    elif isinstance(val, (tuple, list, set)) and len(val) < 20:
+        optional_string = str(val)
+    elif isinstance(val, RaisesExc):
+        optional_string = " or ".join([f"{exc.__name__}" for exc in val.expected_exceptions])
+    elif isinstance(val, nullcontext):
+        optional_string = "does_not_raise()"
+
+    if isinstance(optional_string, str) and len(optional_string) > MAX_SIZE:
+        optional_string = optional_string[: MAX_SIZE - 3] + "+++"  # Can't use dots with pytest
+
+    return optional_string

tests/device.py

@@ -0,0 +1,16 @@
+import os
+
+import torch
+
+try:
+    _device_str = os.environ["PYTEST_TORCH_DEVICE"]
+except KeyError:
+    _device_str = "cpu"  # Default to cpu if environment variable not set
+
+if _device_str != "cuda:0" and _device_str != "cpu":
+    raise ValueError(f"Invalid value of environment variable PYTEST_TORCH_DEVICE: {_device_str}")
+
+if _device_str == "cuda:0" and not torch.cuda.is_available():
+    raise ValueError('Requested device "cuda:0" but cuda is not available.')
+
+DEVICE = torch.device(_device_str)

tests/speed/autogram/grad_vs_jac_vs_gram.py

@@ -2,7 +2,7 @@
 import time
 
 import torch
-from unit.conftest import DEVICE
+from device import DEVICE
 from utils.architectures import (
     AlexNet,
     Cifar10Model,

tests/unit/aggregation/_inputs.py

@@ -1,5 +1,5 @@
 import torch
-from unit.conftest import DEVICE
+from device import DEVICE
 from utils.tensors import zeros_
 
 from ._matrix_samplers import NonWeakSampler, NormalSampler, StrictlyWeakSampler, StrongSampler

tests/unit/autojac/_transform/test_aggregate.py

@@ -1,8 +1,8 @@
 import math
 
 import torch
+from device import DEVICE
 from pytest import mark, raises
-from unit.conftest import DEVICE
 from utils.dict_assertions import assert_tensor_dicts_are_close
 from utils.tensors import rand_, tensor_, zeros_
 

tests/unit/autojac/test_utils.py

@@ -1,6 +1,6 @@
+from device import DEVICE
 from pytest import mark, raises
 from torch.nn import Linear, MSELoss, ReLU, Sequential
-from unit.conftest import DEVICE
 from utils.tensors import randn_, tensor_
 
 from torchjd.autojac._utils import get_leaf_tensors

tests/utils/architectures.py

@@ -2,10 +2,10 @@
 
 import torch
 import torchvision
+from device import DEVICE
 from torch import Tensor, nn
 from torch.nn import Flatten, ReLU
 from torch.utils._pytree import PyTree
-from unit.conftest import DEVICE
 
 
 class ModuleFactory:

tests/utils/tensors.py

@@ -1,8 +1,8 @@
 from functools import partial
 
 import torch
+from device import DEVICE
 from torch.utils._pytree import PyTree, tree_map
-from unit.conftest import DEVICE
 
 # Curried calls to torch functions that require a device so that we automatically fix the device
 # for code written in the tests, while not affecting code written in src (what