diff --git a/tests/async_sleep.py b/tests/async_sleep.py new file mode 100644 index 00000000..ce700003 --- /dev/null +++ b/tests/async_sleep.py @@ -0,0 +1,23 @@ +# 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 asyncio + + +async def sleep(secs): + """Tests can use this function to sleep.""" + await asyncio.sleep(secs) diff --git a/tests/sleep.py b/tests/sleep.py new file mode 100644 index 00000000..b45d71b6 --- /dev/null +++ b/tests/sleep.py @@ -0,0 +1,23 @@ +# 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 time + + +def sleep(secs): + """Tests can use this function to sleep.""" + time.sleep(secs) diff --git a/tests/test_integration/test_examples/_async/test_vectors.py b/tests/test_integration/test_examples/_async/test_vectors.py index 0d634ddf..effb48cf 100644 --- a/tests/test_integration/test_examples/_async/test_vectors.py +++ b/tests/test_integration/test_examples/_async/test_vectors.py @@ -15,20 +15,38 @@ # specific language governing permissions and limitations # under the License. +from hashlib import md5 from unittest import SkipTest import pytest -from ..async_examples.vectors import create, search +from tests.async_sleep import sleep + +from ..async_examples import vectors @pytest.mark.asyncio -async def test_vector_search(async_write_client, es_version): +async def test_vector_search(async_write_client, es_version, mocker): # this test only runs on Elasticsearch >= 8.11 because the example uses - # a dense vector without giving them an explicit size + # a dense vector without specifying an explicit size if es_version < (8, 11): raise SkipTest("This test requires Elasticsearch 8.11 or newer") - await create() - results = await (await search("work from home")).execute() - assert results[0].name == "Work From Home Policy" + class MockModel: + def __init__(self, model): + pass + + def encode(self, text): + vector = [int(ch) for ch in md5(text.encode()).digest()] + total = sum(vector) + return [float(v) / total for v in vector] + + mocker.patch.object(vectors, "SentenceTransformer", new=MockModel) + + await vectors.create() + for i in range(10): + results = await (await vectors.search("Welcome to our team!")).execute() + if len(results.hits) > 0: + break + await sleep(0.1) + assert results[0].name == "New Employee Onboarding Guide" diff --git a/tests/test_integration/test_examples/_sync/test_vectors.py b/tests/test_integration/test_examples/_sync/test_vectors.py index 37aef147..395aa408 100644 --- a/tests/test_integration/test_examples/_sync/test_vectors.py +++ b/tests/test_integration/test_examples/_sync/test_vectors.py @@ -15,20 +15,38 @@ # specific language governing permissions and limitations # under the License. +from hashlib import md5 from unittest import SkipTest import pytest -from ..examples.vectors import create, search +from tests.sleep import sleep + +from ..examples import vectors @pytest.mark.sync -def test_vector_search(write_client, es_version): +def test_vector_search(write_client, es_version, mocker): # this test only runs on Elasticsearch >= 8.11 because the example uses - # a dense vector without giving them an explicit size + # a dense vector without specifying an explicit size if es_version < (8, 11): raise SkipTest("This test requires Elasticsearch 8.11 or newer") - create() - results = (search("work from home")).execute() - assert results[0].name == "Work From Home Policy" + class MockModel: + def __init__(self, model): + pass + + def encode(self, text): + vector = [int(ch) for ch in md5(text.encode()).digest()] + total = sum(vector) + return [float(v) / total for v in vector] + + mocker.patch.object(vectors, "SentenceTransformer", new=MockModel) + + vectors.create() + for i in range(10): + results = (vectors.search("Welcome to our team!")).execute() + if len(results.hits) > 0: + break + sleep(0.1) + assert results[0].name == "New Employee Onboarding Guide" diff --git a/utils/run-unasync.py b/utils/run-unasync.py index db27edcf..aaaba069 100644 --- a/utils/run-unasync.py +++ b/utils/run-unasync.py @@ -69,6 +69,7 @@ def main(check=False): "async_write_client": "write_client", "async_pull_request": "pull_request", "async_examples": "examples", + "async_sleep": "sleep", "assert_awaited_once_with": "assert_called_once_with", "pytest_asyncio": "pytest", }