From 82d4800a61e281d6b03a8b1f7332b38cc9eee6a2 Mon Sep 17 00:00:00 2001
From: Oleksii Hrechykha <alexeee6@gmail.com>
Date: Fri, 19 Jul 2024 08:20:20 +0300
Subject: [PATCH] adding unique_inverse tests

---
 tests/test_unique_inverse_55.py               | 45 +++++++++++++++++++
 ...values1d.py => test_unique_values1d_55.py} |  0
 tests/unique_counts.py                        |  7 ++-
 tests/unique_inverse.py                       | 10 +++--
 tests/unique_values.py                        |  2 +-
 5 files changed, 59 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_unique_inverse_55.py
 rename tests/{test_unique_values1d.py => test_unique_values1d_55.py} (100%)

diff --git a/tests/test_unique_inverse_55.py b/tests/test_unique_inverse_55.py
new file mode 100644
index 0000000..b337d61
--- /dev/null
+++ b/tests/test_unique_inverse_55.py
@@ -0,0 +1,45 @@
+from unique_inverse import unique_inverse
+import awkward as ak
+import ragged
+
+def test_can_take_none():
+    assert unique_inverse(None) is None
+
+def test_can_take_list():
+    assert unique_inverse([1,2,4,3,4,5,6,20]) is None
+
+def test_can_take_empty_arr():
+    arr = ragged.array([])
+    expected_unique_values = ragged.array([])
+    expected_unique_indices = ragged.array([])
+    unique_values, unique_indices = unique_inverse(arr)
+    assert ak.to_list(unique_values) == ak.to_list(expected_unique_values)
+    assert ak.to_list(unique_indices) == ak.to_list(expected_unique_indices)
+
+def test_can_take_simple_array():
+    arr = ragged.array([1, 2, 2, 3, 3, 3, 4, 4, 4, 4])
+    expected_unique_values = ragged.array([1, 2, 3, 4])
+    expected_unique_indices = ragged.array([0, 1, 3, 6])
+    unique_values, unique_indices = unique_inverse(arr)
+    assert ak.to_list(unique_values) == ak.to_list(expected_unique_values)
+    assert ak.to_list(unique_indices) == ak.to_list(expected_unique_indices)
+    
+def test_can_take_normal_array():
+    arr = ragged.array([[1, 2, 2], [3], [3, 3], [4, 4, 4], [4]])
+    expected_unique_values = ragged.array([1, 2, 3, 4])
+    expected_unique_indices = ragged.array([0, 1, 3, 6])
+    unique_values, unique_indices = unique_inverse(arr)
+    assert ak.to_list(unique_values) == ak.to_list(expected_unique_values)
+    assert ak.to_list(unique_indices) == ak.to_list(expected_unique_indices)
+       
+def test_can_take_awkward():
+    arr = ragged.array([[[1, 2, 2], 3, 3], [3, 4, 4, 4, 4]])
+    assert unique_inverse(arr) is None
+
+def test_can_take_scalar():
+    arr = ragged.array([5])
+    expected_unique_values = ragged.array([5])
+    expected_unique_indices = ragged.array([0])
+    unique_values, unique_indices = unique_inverse(arr)
+    assert ak.to_list(unique_values) == ak.to_list(expected_unique_values)
+    assert ak.to_list(unique_indices) == ak.to_list(expected_unique_indices)
\ No newline at end of file
diff --git a/tests/test_unique_values1d.py b/tests/test_unique_values1d_55.py
similarity index 100%
rename from tests/test_unique_values1d.py
rename to tests/test_unique_values1d_55.py
diff --git a/tests/unique_counts.py b/tests/unique_counts.py
index 89690e2..079648b 100644
--- a/tests/unique_counts.py
+++ b/tests/unique_counts.py
@@ -1,7 +1,12 @@
+import awkward as ak
+import ragged
+import numpy as np
+
+
 def unique_counts(arr):
     if not isinstance(arr, ragged.array):
         print("Input is not a ragged array")
-
+        return None
     if len(arr)==1:
         return [(arr[0], 1)]
 
diff --git a/tests/unique_inverse.py b/tests/unique_inverse.py
index 9f951b9..6e6d2d4 100644
--- a/tests/unique_inverse.py
+++ b/tests/unique_inverse.py
@@ -1,11 +1,15 @@
+import awkward as ak
+import ragged
+import numpy as np
+
 def unique_inverse(arr):
     if not isinstance(arr, ragged.array):
         print("Input is not a ragged array")
-        
+        return None
     if len(arr) == 1:
-        return arr, np.array([0])
+        return (arr, np.array([0]))
         
-    arr_list=ak.ravel(arr)
+    arr_list=ak.ravel(arr._impl)
     arr_np = ak.to_numpy(arr_list)
     unique_elements, first_indices = np.unique(arr_np, return_index=True)
     
diff --git a/tests/unique_values.py b/tests/unique_values.py
index cee54ff..a78fec7 100644
--- a/tests/unique_values.py
+++ b/tests/unique_values.py
@@ -9,7 +9,7 @@ def unique_values(arr):
     if len(arr)==1:
         return arr
 
-    arr_flat=ak.ravel(arr)
+    arr_flat=ak.ravel(arr._impl)
     arr_np = ak.to_numpy(arr_flat)
     unique_np = np.unique(arr_np)
     unique_arr = ak.from_numpy(unique_np)