From 0611fdcb3d93838e4771f2bd1f7e783e684325d0 Mon Sep 17 00:00:00 2001
From: Albert Villanova del Moral
 <8515462+albertvillanova@users.noreply.github.com>
Date: Fri, 2 Aug 2024 14:49:46 +0200
Subject: [PATCH] Test generate_from_arrow_type with primitive/nested data
 types

---
 tests/features/test_features.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/tests/features/test_features.py b/tests/features/test_features.py
index 2518c9bac3a..266c2c87066 100644
--- a/tests/features/test_features.py
+++ b/tests/features/test_features.py
@@ -18,6 +18,7 @@
     _check_if_features_can_be_aligned,
     cast_to_python_objects,
     encode_nested_example,
+    generate_from_arrow_type,
     generate_from_dict,
     get_nested_type,
     string_to_arrow,
@@ -804,3 +805,27 @@ def test_get_nested_type_with_list_feature(
     feature = list_feature(scalar_feature)
     arrow_data_type = get_nested_type(feature)
     assert arrow_data_type == expected_arrow_nested_data_type(expected_arrow_primitive_data_type())
+
+
+@pytest.mark.parametrize(
+    "arrow_primitive_data_type, expected_feature", [(pa.int32, Value("int32")), (pa.string, Value("string"))]
+)
+def test_generate_from_arrow_type_with_arrow_primitive_data_type(arrow_primitive_data_type, expected_feature):
+    arrow_data_type = arrow_primitive_data_type()
+    feature = generate_from_arrow_type(arrow_data_type)
+    assert feature == expected_feature
+
+
+@pytest.mark.parametrize(
+    "arrow_primitive_data_type, expected_scalar_feature", [(pa.int32, Value("int32")), (pa.string, Value("string"))]
+)
+@pytest.mark.parametrize(
+    "arrow_nested_data_type, expected_list_feature", [(pa.list_, Sequence), (pa.large_list, LargeList)]
+)
+def test_generate_from_arrow_type_with_arrow_nested_data_type(
+    arrow_nested_data_type, expected_list_feature, arrow_primitive_data_type, expected_scalar_feature
+):
+    arrow_data_type = arrow_nested_data_type(arrow_primitive_data_type())
+    feature = generate_from_arrow_type(arrow_data_type)
+    expected_feature = expected_list_feature(expected_scalar_feature)
+    assert feature == expected_feature