Increasing test-coverage

Author: marcromeyn

merlin/models/torch/batch.py

@@ -281,7 +281,6 @@ def target(self, name: str = "default") -> torch.Tensor:
     def __bool__(self) -> bool:
         return bool(self.features)
 
-<<<<<<< HEAD
     def device(self) -> torch.device:
         """Retrieves the device of the tensors in the Batch object.
 
@@ -303,14 +302,6 @@ def sample_batch(
     batch_size: Optional[int] = None,
     shuffle: Optional[bool] = False,
 ) -> Batch:
-=======
-
-def sample_batch(
-    dataset_or_loader: Union[Dataset, Loader],
-    batch_size: Optional[int] = None,
-    shuffle: Optional[bool] = False,
-) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
->>>>>>> d00e67a4 (Adding sample_batch & sample_features)
     """Util function to generate a batch of input tensors from a merlin.io.Dataset instance
 
     Parameters
@@ -330,6 +321,7 @@ def sample_batch(
         dictionary of target tensors.
     """
 
+<<<<<<< HEAD
 <<<<<<< HEAD
     if isinstance(data, Dataset):
         if not batch_size:
@@ -340,6 +332,11 @@ def sample_batch(
     else:
         raise ValueError(f"Expected Dataset or Loader instance, got: {data}")
 =======
+=======
+    if not isinstance(dataset_or_loader, (Dataset, Loader)):
+        raise ValueError(f"Expected Dataset or Loader instance, got {dataset_or_loader}")
+
+>>>>>>> 519159b5 (Increasing test-coverage)
     if isinstance(dataset_or_loader, Dataset):
         if not batch_size:
             raise ValueError("Either use 'Loader' or specify 'batch_size'")

merlin/models/torch/router.py

@@ -75,14 +75,13 @@ def add_route(
 
         routing_module = self.selectable.select(selection)
         if module is not None:
-            branch = Block(routing_module)
-            if isinstance(module, ParallelBlock):
-                branch = module.prepend(routing_module)
-
             if hasattr(module, "setup_schema"):
                 module.setup_schema(routing_module.schema)
 
-            branch.append(module)
+            if isinstance(module, ParallelBlock):
+                branch = module.prepend(routing_module)
+            else:
+                branch = Block(routing_module, module)
         else:
             branch = routing_module
 

tests/unit/torch/utils/test_selection_utils.py

@@ -1,6 +1,6 @@
 import pytest
 
-from merlin.models.torch.utils.selection_utils import select_schema, selection_name
+from merlin.models.torch.utils.selection_utils import Selectable, select_schema, selection_name
 from merlin.schema import ColumnSchema, Schema, Tags
 
 
@@ -70,3 +70,16 @@ def test_select_column(self):
 
         assert selection_name(column) == column.name
         assert selection_name(ColumnSchema("user_id")) == column.name
+
+    def test_exception(self):
+        with pytest.raises(ValueError, match="is not valid"):
+            selection_name(1)
+
+
+class TestSelectable:
+    def test_exception(self):
+        selectable = Selectable()
+
+        assert hasattr(selectable, "setup_schema")
+        with pytest.raises(NotImplementedError):
+            selectable.select(1)