Add test_dataloaders to test method

pytorch_lightning/trainer/trainer.py

@@ -612,8 +612,7 @@ def fit(
             self,
             model: LightningModule,
             train_dataloader: Optional[DataLoader] = None,
-            val_dataloaders: Optional[DataLoader] = None,
-            test_dataloaders: Optional[DataLoader] = None
+            val_dataloaders: Optional[DataLoader] = None
     ):
         r"""
         Runs the full optimization routine.
@@ -629,10 +628,6 @@ def fit(
                 Pytorch Dataloader or a list of them, specifying validation samples.
                 If the model has a predefined val_dataloaders method this will be skipped
 
-            test_dataloaders: Either a single
-                Pytorch Dataloader or a list of them, specifying validation samples.
-                If the model has a predefined test_dataloaders method this will be skipped
-
         Example::
 
             # Option 1,
@@ -646,11 +641,10 @@ def fit(
             # Option 2
             # in production cases we might want to pass different datasets to the same model
             # Recommended for PRODUCTION SYSTEMS
-            train, val, test = DataLoader(...), DataLoader(...), DataLoader(...)
+            train, val = DataLoader(...), DataLoader(...)
             trainer = Trainer()
             model = LightningModule()
-            trainer.fit(model, train_dataloader=train,
-                        val_dataloader=val, test_dataloader=test)
+            trainer.fit(model, train_dataloader=train, val_dataloader=val)
 
             # Option 1 & 2 can be mixed, for example the training set can be
             # defined as part of the model, and validation/test can then be
@@ -662,7 +656,7 @@ def fit(
         self.copy_trainer_model_properties(model)
 
         # set up the passed in dataloaders (if needed)
-        self.__attach_dataloaders(model, train_dataloader, val_dataloaders, test_dataloaders)
+        self.__attach_dataloaders(model, train_dataloader, val_dataloaders)
 
         # check that model is configured correctly
         self.check_model_configuration(model)
@@ -747,7 +741,7 @@ def __set_random_port(self):
             default_port = random.randint(10000, 19000)
             os.environ['MASTER_PORT'] = str(default_port)
 
-    def __attach_dataloaders(self, model, train_dataloader, val_dataloaders, test_dataloaders):
+    def __attach_dataloaders(self, model, train_dataloader=None, val_dataloaders=None, test_dataloaders=None):
         # when dataloader is passed via fit, patch the train_dataloader
         # functions to overwrite with these implementations
         if train_dataloader is not None:
@@ -863,14 +857,17 @@ def run_pretrain_routine(self, model: LightningModule):
         # CORE TRAINING LOOP
         self.train()
 
-    def test(self, model: Optional[LightningModule] = None):
+    def test(self, model: Optional[LightningModule] = None, test_dataloaders: Optional[DataLoader] = None):
         r"""
 
         Separates from fit to make sure you never run on your test set until you want to.
 
         Args:
             model: The model to test.
 
+            test_dataloaders: Either a single
+                Pytorch Dataloader or a list of them, specifying validation samples.
+
         Example::
 
             # Option 1
@@ -889,6 +886,13 @@ def test(self, model: Optional[LightningModule] = None):
         """
 
         self.testing = True
+
+        if test_dataloaders is not None:
+            if model is not None:
+                self.__attach_dataloaders(model, test_dataloaders=test_dataloaders)
+            else:
+                self.__attach_dataloaders(self.model, test_dataloaders=test_dataloaders)
+
         if model is not None:
             self.model = model
             self.fit(model)