From d70f5e730a005fa561348c0653f310e512c6095b Mon Sep 17 00:00:00 2001
From: Francesco Andreuzzi <andreuzzi.francesco@gmail.com>
Date: Thu, 8 Dec 2022 22:37:06 +0100
Subject: [PATCH] test on merge_tensors

---
 tests/test_utils.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)
 create mode 100644 tests/test_utils.py

diff --git a/tests/test_utils.py b/tests/test_utils.py
new file mode 100644
index 0000000..7ed795f
--- /dev/null
+++ b/tests/test_utils.py
@@ -0,0 +1,15 @@
+import torch
+
+from pina.utils import merge_tensors
+from pina.label_tensor import LabelTensor
+
+def test_merge_tensors():
+    tensor1 = LabelTensor(torch.rand((20, 3)), ['a', 'b', 'c'])
+    tensor2 = LabelTensor(torch.zeros((20, 3)), ['d', 'e', 'f'])
+    tensor3 = LabelTensor(torch.ones((30, 3)), ['g', 'h', 'i'])
+
+    merged_tensor = merge_tensors((tensor1, tensor2, tensor3))
+    assert tuple(merged_tensor.labels) == ('a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i')
+    assert merged_tensor.shape == (20*20*30, 9)
+    assert torch.all(merged_tensor.extract(('d', 'e', 'f')) == 0)
+    assert torch.all(merged_tensor.extract(('g', 'h', 'i')) == 1)