From d692c70e571f1311d04ac7b124c89b73d73606d4 Mon Sep 17 00:00:00 2001 From: danigm Date: Wed, 8 Jan 2025 12:40:33 +0100 Subject: [PATCH] Make pyarrow dependency optional for tests (#2733) --- .../test_client/test_deprecated_options.py | 39 ++++++++------- .../test_client/test_serializers.py | 49 ++++++++----------- test_elasticsearch/test_serializer.py | 16 +++--- 3 files changed, 49 insertions(+), 55 deletions(-) diff --git a/test_elasticsearch/test_client/test_deprecated_options.py b/test_elasticsearch/test_client/test_deprecated_options.py index 810e75cf4..3b4e0b9ed 100644 --- a/test_elasticsearch/test_client/test_deprecated_options.py +++ b/test_elasticsearch/test_client/test_deprecated_options.py @@ -21,6 +21,23 @@ from elasticsearch import Elasticsearch, JsonSerializer +EXPECTED_SERIALIZERS = { + "application/vnd.mapbox-vector-tile", + "application/x-ndjson", + "application/json", + "text/*", + "application/vnd.elasticsearch+json", + "application/vnd.elasticsearch+x-ndjson", +} + + +try: + import pyarrow as pa + + EXPECTED_SERIALIZERS.add("application/vnd.apache.arrow.stream") +except ImportError: + pa = None + def test_sniff_on_connection_fail(): with warnings.catch_warnings(record=True) as w: @@ -129,15 +146,7 @@ class CustomSerializer(JsonSerializer): client.transport.serializers.get_serializer("application/json"), CustomSerializer, ) - assert set(client.transport.serializers.serializers.keys()) == { - "application/vnd.mapbox-vector-tile", - "application/x-ndjson", - "application/json", - "text/*", - "application/vnd.apache.arrow.stream", - "application/vnd.elasticsearch+json", - "application/vnd.elasticsearch+x-ndjson", - } + assert set(client.transport.serializers.serializers.keys()) == EXPECTED_SERIALIZERS client = Elasticsearch( "http://localhost:9200", @@ -150,13 +159,5 @@ class CustomSerializer(JsonSerializer): client.transport.serializers.get_serializer("application/json"), CustomSerializer, ) - assert set(client.transport.serializers.serializers.keys()) == { - "application/vnd.mapbox-vector-tile", - "application/x-ndjson", - "application/json", - "text/*", - "application/vnd.apache.arrow.stream", - "application/vnd.elasticsearch+json", - "application/vnd.elasticsearch+x-ndjson", - "application/cbor", - } + expected = EXPECTED_SERIALIZERS | {"application/cbor"} + assert set(client.transport.serializers.serializers.keys()) == expected diff --git a/test_elasticsearch/test_client/test_serializers.py b/test_elasticsearch/test_client/test_serializers.py index 9d13386ed..986160b92 100644 --- a/test_elasticsearch/test_client/test_serializers.py +++ b/test_elasticsearch/test_client/test_serializers.py @@ -20,6 +20,23 @@ from elasticsearch import Elasticsearch from test_elasticsearch.test_cases import DummyTransportTestCase +EXPECTED_SERIALIZERS = { + "application/json", + "text/*", + "application/x-ndjson", + "application/vnd.mapbox-vector-tile", + "application/vnd.elasticsearch+json", + "application/vnd.elasticsearch+x-ndjson", +} + + +try: + import pyarrow as pa + + EXPECTED_SERIALIZERS.add("application/vnd.apache.arrow.stream") +except ImportError: + pa = None + class TestSerializers(DummyTransportTestCase): def test_compat_mode_on_by_default(self): @@ -90,16 +107,8 @@ class CustomSerializer: "https://localhost:9200", serializers={f"application/{mime_subtype}": ser} ) serializers = client.transport.serializers.serializers - assert set(serializers.keys()) == { - "application/json", - "text/*", - "application/x-ndjson", - "application/vnd.apache.arrow.stream", - "application/vnd.mapbox-vector-tile", - "application/vnd.elasticsearch+json", - "application/vnd.elasticsearch+x-ndjson", - } + assert set(serializers.keys()) == EXPECTED_SERIALIZERS assert serializers[f"application/{mime_subtype}"] is ser assert serializers[f"application/vnd.elasticsearch+{mime_subtype}"] is ser @@ -118,16 +127,7 @@ class CustomSerializer: }, ) serializers = client.transport.serializers.serializers - assert set(serializers.keys()) == { - "application/json", - "text/*", - "application/x-ndjson", - "application/vnd.apache.arrow.stream", - "application/vnd.mapbox-vector-tile", - "application/vnd.elasticsearch+json", - "application/vnd.elasticsearch+x-ndjson", - } - + assert set(serializers.keys()) == EXPECTED_SERIALIZERS assert serializers[f"application/{mime_subtype}"] is ser1 assert serializers[f"application/vnd.elasticsearch+{mime_subtype}"] is ser2 @@ -138,15 +138,6 @@ class CustomSerializer: ser = CustomSerializer() client = Elasticsearch("https://localhost:9200", serializer=ser) serializers = client.transport.serializers.serializers - assert set(serializers.keys()) == { - "application/json", - "text/*", - "application/x-ndjson", - "application/vnd.apache.arrow.stream", - "application/vnd.mapbox-vector-tile", - "application/vnd.elasticsearch+json", - "application/vnd.elasticsearch+x-ndjson", - } - + assert set(serializers.keys()) == EXPECTED_SERIALIZERS assert serializers["application/json"] is ser assert serializers["application/vnd.elasticsearch+json"] is ser diff --git a/test_elasticsearch/test_serializer.py b/test_elasticsearch/test_serializer.py index 02723e8f4..21c6b94b5 100644 --- a/test_elasticsearch/test_serializer.py +++ b/test_elasticsearch/test_serializer.py @@ -19,9 +19,15 @@ from datetime import datetime from decimal import Decimal -import pyarrow as pa import pytest +try: + import pyarrow as pa + + from elasticsearch.serializer import PyArrowSerializer +except ImportError: + pa = None + try: import numpy as np import pandas as pd @@ -32,12 +38,7 @@ from elasticsearch import Elasticsearch from elasticsearch.exceptions import SerializationError -from elasticsearch.serializer import ( - JSONSerializer, - OrjsonSerializer, - PyArrowSerializer, - TextSerializer, -) +from elasticsearch.serializer import JSONSerializer, OrjsonSerializer, TextSerializer requires_numpy_and_pandas = pytest.mark.skipif( np is None or pd is None, reason="Test requires numpy and pandas to be available" @@ -163,6 +164,7 @@ def test_serializes_pandas_category(json_serializer): assert b'{"d":[1,2,3]}' == json_serializer.dumps({"d": cat}) +@pytest.mark.skipif(pa is None, reason="Test requires pyarrow to be available") def test_pyarrow_loads(): data = [ pa.array([1, 2, 3, 4]),