1+ import  warnings 
2+ from  urllib .error  import  URLError 
3+ 
14import  pytest 
25
36from  torch_uncertainty .datamodules .classification  import  (
@@ -13,27 +16,30 @@ class TestHTRU2DataModule:
1316    """Testing the HTRU2DataModule datamodule class.""" 
1417
1518    def  test_htru2 (self ) ->  None :
16-         dm  =  HTRU2DataModule (root = "./data/" , batch_size = 128 )
19+         try :
20+             dm  =  HTRU2DataModule (root = "./data/" , batch_size = 128 )
1721
18-         dm .prepare_data ()
19-         dm .setup ()
22+              dm .prepare_data ()
23+              dm .setup ()
2024
21-         dm .train_dataloader ()
22-         dm .val_dataloader ()
23-         dm .test_dataloader ()
25+              dm .train_dataloader ()
26+              dm .val_dataloader ()
27+              dm .test_dataloader ()
2428
25-         dm .setup ("test" )
26-         dm .test_dataloader ()
29+              dm .setup ("test" )
30+              dm .test_dataloader ()
2731
28-         dm  =  HTRU2DataModule (root = "./data/" , batch_size = 128 , val_split = 0.1 )
32+              dm  =  HTRU2DataModule (root = "./data/" , batch_size = 128 , val_split = 0.1 )
2933
30-         dm .prepare_data ()
31-         dm .setup ()
34+              dm .prepare_data ()
35+              dm .setup ()
3236
33-         with  pytest .raises (ValueError ):
34-             dm .setup ("other" )
37+              with  pytest .raises (ValueError ):
38+                  dm .setup ("other" )
3539
36-         dm  =  BankMarketingDataModule (root = "./data/" , batch_size = 128 )
37-         dm  =  DOTA2GamesDataModule (root = "./data/" , batch_size = 128 )
38-         dm  =  OnlineShoppersDataModule (root = "./data/" , batch_size = 128 )
39-         dm  =  SpamBaseDataModule (root = "./data/" , batch_size = 128 )
40+             dm  =  BankMarketingDataModule (root = "./data/" , batch_size = 128 )
41+             dm  =  DOTA2GamesDataModule (root = "./data/" , batch_size = 128 )
42+             dm  =  OnlineShoppersDataModule (root = "./data/" , batch_size = 128 )
43+             dm  =  SpamBaseDataModule (root = "./data/" , batch_size = 128 )
44+         except  URLError  as  e :
45+             warnings .warn (f"Data download failed due to network error: { e }  " , stacklevel = 2 )
0 commit comments