Skip to content

Commit

Permalink
use elasticsearch_url fixture; create conftest.py
Browse files Browse the repository at this point in the history
  • Loading branch information
maxjakob committed Apr 22, 2024
1 parent 7647961 commit d4a84c1
Show file tree
Hide file tree
Showing 9 changed files with 207 additions and 273 deletions.
74 changes: 0 additions & 74 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,77 +94,3 @@
"orjson": ["orjson>=3"],
},
)

vectorstore_package_name = "elasticsearch[vectorstore]"
base_dir = abspath(dirname(__file__))

with open(join(base_dir, package_name, "_version.py")) as f:
package_version = re.search(
r"__versionstr__\s+=\s+[\"\']([^\"\']+)[\"\']", f.read()
).group(1)

with open(join(base_dir, "README.rst")) as f:
# Remove reST raw directive from README as they're not allowed on PyPI
# Those blocks start with a newline and continue until the next newline
mode = None
lines = []
for line in f:
if line.startswith(".. raw::"):
mode = "ignore_nl"
elif line == "\n":
mode = "wait_nl" if mode == "ignore_nl" else None
if mode is None:
lines.append(line)

long_description = "".join(lines)


packages = [
package
for package in find_packages(where=".", exclude=("test_elasticsearch*",))
if package == package_name or package.startswith(package_name + ".")
]

setup(
name=package_name,
description="Python client for Elasticsearch",
license="Apache-2.0",
url="https://github.com/elastic/elasticsearch-py",
long_description=long_description,
long_description_content_type="text/x-rst",
version=package_version,
author="Elastic Client Library Maintainers",
author_email="[email protected]",
project_urls={
"Documentation": "https://elasticsearch-py.readthedocs.io",
"Source Code": "https://github.com/elastic/elasticsearch-py",
"Issue Tracker": "https://github.com/elastic/elasticsearch-py/issues",
},
packages=packages,
package_data={"elasticsearch": ["py.typed", "*.pyi"]},
include_package_data=True,
zip_safe=False,
classifiers=[
"Development Status :: 5 - Production/Stable",
"License :: OSI Approved :: Apache Software License",
"Intended Audience :: Developers",
"Operating System :: OS Independent",
"Programming Language :: Python",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.7",
"Programming Language :: Python :: 3.8",
"Programming Language :: Python :: 3.9",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: Implementation :: CPython",
"Programming Language :: Python :: Implementation :: PyPy",
],
python_requires=">=3.7",
install_requires=["elastic-transport>=8.13,<9"],
extras_require={
"requests": ["requests>=2.4.0, <3.0.0"],
"async": ["aiohttp>=3,<4"],
"orjson": ["orjson>=3"],
},
)
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,10 @@
# specific language governing permissions and limitations
# under the License.

import os
from typing import Any, AsyncIterator, Dict, List, Optional
from typing import Any, Dict, List

from elastic_transport import AsyncTransport

from elasticsearch import AsyncElasticsearch
from elasticsearch.vectorstore._async.embedding_service import AsyncEmbeddingService


Expand Down Expand Up @@ -87,69 +85,3 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
async def perform_request(self, *args, **kwargs): # type: ignore
self.requests.append(kwargs)
return await super().perform_request(*args, **kwargs)


def create_es_client(
es_params: Optional[Dict[str, str]] = None, es_kwargs: Dict = {}
) -> AsyncElasticsearch:
if es_params is None:
es_params = read_env()
if not es_kwargs:
es_kwargs = {}

if "es_cloud_id" in es_params:
return AsyncElasticsearch(
cloud_id=es_params["es_cloud_id"],
api_key=es_params["es_api_key"],
**es_kwargs,
)
return AsyncElasticsearch(hosts=[es_params["es_url"]], **es_kwargs)


def create_requests_saving_client() -> AsyncElasticsearch:
return create_es_client(es_kwargs={"transport_class": AsyncRequestSavingTransport})


async def es_client_fixture() -> AsyncIterator[AsyncElasticsearch]:
params = read_env()
client = create_es_client(params)

yield client

# clear indices
await clear_test_indices(client)

# clear all test pipelines
try:
response = await client.ingest.get_pipeline(id="test_*,*_sparse_embedding")

for pipeline_id, _ in response.items():
try:
await client.ingest.delete_pipeline(id=pipeline_id)
print(f"Deleted pipeline: {pipeline_id}") # noqa: T201
except Exception as e:
print(f"Pipeline error: {e}") # noqa: T201

except Exception:
pass
finally:
await client.close()


async def clear_test_indices(client: AsyncElasticsearch) -> None:
response = await client.indices.get(index="_all")
index_names = response.keys()
for index_name in index_names:
if index_name.startswith("test_"):
await client.indices.delete(index=index_name)
await client.indices.refresh(index="_all")


def read_env() -> Dict:
url = os.environ.get("ES_URL", "http://localhost:9200")
cloud_id = os.environ.get("ES_CLOUD_ID")
api_key = os.environ.get("ES_API_KEY")

if cloud_id:
return {"es_cloud_id": cloud_id, "es_api_key": api_key}
return {"es_url": url}
103 changes: 103 additions & 0 deletions test_elasticsearch/test_server/test_vectorstore/_async/conftest.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
# Licensed to Elasticsearch B.V. under one or more contributor
# license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright
# ownership. Elasticsearch B.V. licenses this file to you under
# the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.

import os
import uuid
from typing import AsyncIterator, Dict

import pytest
import pytest_asyncio

from elasticsearch import AsyncElasticsearch

from ._test_utils import AsyncRequestSavingTransport


@pytest_asyncio.fixture
async def es_client(elasticsearch_url: str) -> AsyncIterator[AsyncElasticsearch]:
client = _create_es_client(elasticsearch_url)

yield client

# clear indices
await _clear_test_indices(client)

# clear all test pipelines
try:
response = await client.ingest.get_pipeline(id="test_*,*_sparse_embedding")

for pipeline_id, _ in response.items():
try:
await client.ingest.delete_pipeline(id=pipeline_id)
print(f"Deleted pipeline: {pipeline_id}") # noqa: T201
except Exception as e:
print(f"Pipeline error: {e}") # noqa: T201

except Exception:
pass
finally:
await client.close()


@pytest_asyncio.fixture
async def requests_saving_client(
elasticsearch_url: str,
) -> AsyncIterator[AsyncElasticsearch]:
client = _create_es_client(
elasticsearch_url, es_kwargs={"transport_class": AsyncRequestSavingTransport}
)

try:
yield client
finally:
await client.close()


@pytest.fixture(scope="function")
def index_name() -> str:
return f"test_{uuid.uuid4().hex}"


async def _clear_test_indices(client: AsyncElasticsearch) -> None:
response = await client.indices.get(index="_all")
index_names = response.keys()
for index_name in index_names:
if index_name.startswith("test_"):
await client.indices.delete(index=index_name)
await client.indices.refresh(index="_all")


def _create_es_client(
elasticsearch_url: str, es_kwargs: Dict = {}
) -> AsyncElasticsearch:
if not elasticsearch_url:
elasticsearch_url = os.environ.get("ES_URL", "http://localhost:9200")
cloud_id = os.environ.get("ES_CLOUD_ID")
api_key = os.environ.get("ES_API_KEY")

if cloud_id:
es_params = {"es_cloud_id": cloud_id, "es_api_key": api_key}
else:
es_params = {"es_url": elasticsearch_url}

if "es_cloud_id" in es_params:
return AsyncElasticsearch(
cloud_id=es_params["es_cloud_id"],
api_key=es_params["es_api_key"],
**es_kwargs,
)
return AsyncElasticsearch(hosts=[es_params["es_url"]], **es_kwargs)
Original file line number Diff line number Diff line change
Expand Up @@ -27,20 +27,12 @@
AsyncElasticsearchEmbeddings,
)

from ._test_utils import es_client_fixture

# deployed with
# https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-text-emb-vector-search-example.html
MODEL_ID = os.getenv("MODEL_ID", "sentence-transformers__msmarco-minilm-l-12-v3")
NUM_DIMENSIONS = int(os.getenv("NUM_DIMENTIONS", "384"))


@pytest_asyncio.fixture
async def es_client() -> AsyncIterator[AsyncElasticsearch]:
async for x in es_client_fixture():
yield x


@pytest.mark.asyncio
async def test_elasticsearch_embedding_documents(es_client: AsyncElasticsearch) -> None:
"""Test Elasticsearch embedding documents."""
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,10 @@
# under the License.

import logging
import uuid
from functools import partial
from typing import Any, AsyncIterator, List, Optional, Union, cast
from typing import Any, List, Optional, Union, cast

import pytest
import pytest_asyncio

from elasticsearch import AsyncElasticsearch, NotFoundError
from elasticsearch.helpers import BulkIndexError
Expand All @@ -39,8 +37,6 @@
AsyncConsistentFakeEmbeddings,
AsyncFakeEmbeddings,
AsyncRequestSavingTransport,
create_requests_saving_client,
es_client_fixture,
)

logging.basicConfig(level=logging.DEBUG)
Expand All @@ -67,24 +63,6 @@


class TestVectorStore:
@pytest_asyncio.fixture
async def es_client(self) -> AsyncIterator[AsyncElasticsearch]:
async for x in es_client_fixture():
yield x

@pytest_asyncio.fixture
async def requests_saving_client(self) -> AsyncIterator[AsyncElasticsearch]:
client = create_requests_saving_client()
try:
yield client
finally:
await client.close()

@pytest.fixture(scope="function")
def index_name(self) -> str:
"""Return the index name."""
return f"test_{uuid.uuid4().hex}"

@pytest.mark.asyncio
async def test_search_without_metadata(
self, es_client: AsyncElasticsearch, index_name: str
Expand Down
Loading

0 comments on commit d4a84c1

Please sign in to comment.