



import re
import tempfile
from pathlib import Path

import pytest
import yaml

from datasets.utils.readme import ReadMe


# @pytest.fixture
# def example_yaml_structure():

lowercase_	       =  yaml.safe_load(
    '''\\nname: ""\nallow_empty: false\nallow_empty_text: true\nsubsections:\n  - name: "Dataset Card for X" # First-level markdown heading\n    allow_empty: false\n    allow_empty_text: true\n    subsections:\n      - name: "Table of Contents"\n        allow_empty: false\n        allow_empty_text: false\n        subsections: null\n      - name: "Dataset Description"\n        allow_empty: false\n        allow_empty_text: false\n        subsections:\n          - name: "Dataset Summary"\n            allow_empty: false\n            allow_empty_text: false\n            subsections: null\n          - name: "Supported Tasks and Leaderboards"\n            allow_empty: true\n            allow_empty_text: true\n            subsections: null\n          - name: Languages\n            allow_empty: false\n            allow_empty_text: true\n            subsections: null\n'''
)


lowercase_	       =  {
    'name': 'root',
    'text': '',
    'is_empty_text': True,
    'subsections': [
        {
            'name': 'Dataset Card for My Dataset',
            'text': '',
            'is_empty_text': True,
            'subsections': [
                {'name': 'Table of Contents', 'text': 'Some text here.', 'is_empty_text': False, 'subsections': []},
                {
                    'name': 'Dataset Description',
                    'text': 'Some text here.',
                    'is_empty_text': False,
                    'subsections': [
                        {
                            'name': 'Dataset Summary',
                            'text': 'Some text here.',
                            'is_empty_text': False,
                            'subsections': [],
                        },
                        {
                            'name': 'Supported Tasks and Leaderboards',
                            'text': '',
                            'is_empty_text': True,
                            'subsections': [],
                        },
                        {'name': 'Languages', 'text': 'Language Text', 'is_empty_text': False, 'subsections': []},
                    ],
                },
            ],
        }
    ],
}


lowercase_	       =  '\\n---\nlanguage:\n- zh\n- en\n---\n\n# Dataset Card for My Dataset\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\nSome text here.\n### Supported Tasks and Leaderboards\n### Languages\nLanguage Text\n'


lowercase_	       =  '\\n---\nlanguage:\n- zh\n- en\n---\n\n# Dataset Card for My Dataset\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\nSome text here.\n#### Extra Ignored Subsection\n### Supported Tasks and Leaderboards\n### Languages\nLanguage Text\n'

lowercase_	       =  {
    'name': 'root',
    'text': '',
    'is_empty_text': True,
    'subsections': [
        {
            'name': 'Dataset Card for My Dataset',
            'text': '',
            'is_empty_text': True,
            'subsections': [
                {'name': 'Table of Contents', 'text': 'Some text here.', 'is_empty_text': False, 'subsections': []},
                {
                    'name': 'Dataset Description',
                    'text': 'Some text here.',
                    'is_empty_text': False,
                    'subsections': [
                        {
                            'name': 'Dataset Summary',
                            'text': 'Some text here.',
                            'is_empty_text': False,
                            'subsections': [
                                {
                                    'name': 'Extra Ignored Subsection',
                                    'text': '',
                                    'is_empty_text': True,
                                    'subsections': [],
                                }
                            ],
                        },
                        {
                            'name': 'Supported Tasks and Leaderboards',
                            'text': '',
                            'is_empty_text': True,
                            'subsections': [],
                        },
                        {'name': 'Languages', 'text': 'Language Text', 'is_empty_text': False, 'subsections': []},
                    ],
                },
            ],
        }
    ],
}

lowercase_	       =  '\\n---\n---\n# Dataset Card for My Dataset\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\nSome text here.\n### Supported Tasks and Leaderboards\n### Languages\nLanguage Text\n'

lowercase_	       =  (
    'The following issues were found for the README at `{path}`:\n-\tEmpty YAML markers are present in the README.'
)

lowercase_	       =  '\\n# Dataset Card for My Dataset\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\nSome text here.\n### Supported Tasks and Leaderboards\n### Languages\nLanguage Text\n'

lowercase_	       =  (
    'The following issues were found for the README at `{path}`:\n-\tNo YAML markers are present in the README.'
)

lowercase_	       =  '\\n---\n# Dataset Card for My Dataset\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\nSome text here.\n### Supported Tasks and Leaderboards\n### Languages\nLanguage Text\n'

lowercase_	       =  'The following issues were found for the README at `{path}`:\n-\tOnly the start of YAML tags present in the README.'

lowercase_	       =  '\\n---\nlanguage:\n- zh\n- en\n---\n\n# Dataset Card for My Dataset\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\n### Supported Tasks and Leaderboards\n### Languages\nLanguage Text\n'
lowercase_	       =  'The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Dataset Summary` but it is empty.\n-\tExpected some text in section `Dataset Summary` but it is empty (text in subsections are ignored).'


lowercase_	       =  '\\n---\nlanguage:\n- zh\n- en\n---\n\n# Dataset Card for My Dataset\n'
lowercase_	       =  'The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Dataset Card for My Dataset` but it is empty.\n-\tSection `Dataset Card for My Dataset` expected the following subsections: `Table of Contents`, `Dataset Description`. Found \'None\'.'

lowercase_	       =  '\\n---\nlanguage:\n- zh\n- en\n---\n\n# Dataset Card for My Dataset\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\nSome text here.\n### Languages\nLanguage Text\n'

lowercase_	       =  'The following issues were found for the README at `{path}`:\n-\tSection `Dataset Description` is missing subsection: `Supported Tasks and Leaderboards`.'


lowercase_	       =  '\\n---\nlanguage:\n- zh\n- en\n---\n\n# Dataset Card for My Dataset\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\nSome text here.\n### Supported Tasks and Leaderboards\n### Languages\n'

lowercase_	       =  'The following issues were found for the README at `{path}`:\n-\tExpected some content in section `Languages` but it is empty.'

lowercase_	       =  '\\n---\nlanguage:\n- zh\n- en\n---\n\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\nSome text here.\n### Supported Tasks and Leaderboards\n### Languages\nLanguage Text\n'
lowercase_	       =  'The following issues were found for the README at `{path}`:\n-\tThe README has no first-level headings. One heading is expected. Skipping further validation for this README.'

lowercase_	       =  '\\n---\nlanguage:\n- zh\n- en\n---\n\n# Dataset Card for My Dataset\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\nSome text here.\n### Supported Tasks and Leaderboards\n### Languages\nLanguage Text\n# Dataset Card My Dataset\n'

lowercase_	       =  'The following issues were found for the README at `{path}`:\n-\tThe README has several first-level headings: `Dataset Card for My Dataset`, `Dataset Card My Dataset`. Only one heading is expected. Skipping further validation for this README.'

lowercase_	       =  '\\n---\nlanguage:\n- zh\n- en\n---\n\n# Dataset Card My Dataset\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\nSome text here.\n### Supported Tasks and Leaderboards\n### Languages\nLanguage Text\n'

lowercase_	       =  'The following issues were found for the README at `{path}`:\n-\tNo first-level heading starting with `Dataset Card for` found in README. Skipping further validation for this README.'

lowercase_	       =  ''

lowercase_	       =  'The following issues were found for the README at `{path}`:\n-\tThe README has no first-level headings. One heading is expected. Skipping further validation for this README.\n-\tNo YAML markers are present in the README.'

lowercase_	       =  '\\n---\nlanguage:\n- zh\n- en\n---\n\n# Dataset Card for My Dataset\n# Dataset Card for My Dataset\n## Table of Contents\nSome text here.\n## Dataset Description\nSome text here.\n### Dataset Summary\nSome text here.\n### Supported Tasks and Leaderboards\n### Languages\nLanguage Text\n'

lowercase_	       =  'The following issues were found while parsing the README at `{path}`:\n-\tMultiple sections with the same heading `Dataset Card for My Dataset` have been found. Please keep only one of these sections.'






@pytest.mark.parametrize(
    """readme_md, expected_dict"""				,   [
        (README_CORRECT, CORRECT_DICT),
        (README_CORRECT_FOUR_LEVEL, CORRECT_DICT_FOUR_LEVEL),
    ]				,   )
def 					__lowerCAmelCase		(  __lowerCamelCase	:    Dict				,   __lowerCamelCase	:    int	)  ->						Optional[int]:
					assert ReadMe.from_string(A_				,   A_	).to_dict() == expected_dict






@pytest.mark.parametrize(
    """readme_md, expected_error"""				,   [
        (README_NO_YAML, EXPECTED_ERROR_README_NO_YAML),
        (README_EMPTY_YAML, EXPECTED_ERROR_README_EMPTY_YAML),
        (README_INCORRECT_YAML, EXPECTED_ERROR_README_INCORRECT_YAML),
        (README_EMPTY, EXPECTED_ERROR_README_EMPTY),
        (README_NONE_SUBSECTION, EXPECTED_ERROR_README_NONE_SUBSECTION),
        (README_MISSING_FIRST_LEVEL, EXPECTED_ERROR_README_MISSING_FIRST_LEVEL),
        (README_MISSING_SUBSECTION, EXPECTED_ERROR_README_MISSING_SUBSECTION),
        (README_MISSING_TEXT, EXPECTED_ERROR_README_MISSING_TEXT),
        (README_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_WRONG_FIRST_LEVEL),
        (README_MULTIPLE_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_MULTIPLE_WRONG_FIRST_LEVEL),
        (README_MISSING_CONTENT, EXPECTED_ERROR_README_MISSING_CONTENT),
    ]				,   )
def 					__lowerCAmelCase		(  __lowerCamelCase	:    List[Any]				,   __lowerCamelCase	:    Union[str, Any]	)  ->						str:





					with pytest.raises(A_				,   match=re.escape(expected_error.format(path="""root"""	)	)	):
										__lowerCAmelCase	   =ReadMe.from_string(A_				,   A_	)
										readme.validate()






@pytest.mark.parametrize(
    """readme_md, expected_error"""				,   [
        (README_MULTIPLE_SAME_HEADING_1, EXPECTED_ERROR_README_MULTIPLE_SAME_HEADING_1),
    ]				,   )
def 					__lowerCAmelCase		(  __lowerCamelCase	:    List[str]				,   __lowerCamelCase	:    str	)  ->						List[str]:





					with pytest.raises(A_				,   match=re.escape(expected_error.format(path="""root"""	)	)	):
										ReadMe.from_string(A_				,   A_	)






@pytest.mark.parametrize(
    """readme_md,"""				,   [
        (README_MULTIPLE_SAME_HEADING_1),
    ]				,   )
def 					__lowerCAmelCase		(  __lowerCamelCase	:    Optional[Any]	)  ->						List[Any]:
					ReadMe.from_string(A_				,   A_				,   suppress_parsing_errors=A_	)






@pytest.mark.parametrize(
    """readme_md, expected_dict"""				,   [
        (README_CORRECT, CORRECT_DICT),
        (README_CORRECT_FOUR_LEVEL, CORRECT_DICT_FOUR_LEVEL),
    ]				,   )
def 					__lowerCAmelCase		(  __lowerCamelCase	:    int				,   __lowerCamelCase	:    Tuple	)  ->						Optional[Any]:





					with tempfile.TemporaryDirectory() as tmp_dir:
										__lowerCAmelCase	   =Path(A_	) / """README.md"""
										with open(A_				,   """w+"""	) as readme_file:
															readme_file.write(A_	)
										__lowerCAmelCase	   =ReadMe.from_readme(A_				,   A_	).to_dict()
										assert out["name"] == path
										assert out["text"] == ""
										assert out["is_empty_text"]
										assert out["subsections"] == expected_dict["subsections"]






@pytest.mark.parametrize(
    """readme_md, expected_error"""				,   [
        (README_NO_YAML, EXPECTED_ERROR_README_NO_YAML),
        (README_EMPTY_YAML, EXPECTED_ERROR_README_EMPTY_YAML),
        (README_INCORRECT_YAML, EXPECTED_ERROR_README_INCORRECT_YAML),
        (README_EMPTY, EXPECTED_ERROR_README_EMPTY),
        (README_NONE_SUBSECTION, EXPECTED_ERROR_README_NONE_SUBSECTION),
        (README_MISSING_FIRST_LEVEL, EXPECTED_ERROR_README_MISSING_FIRST_LEVEL),
        (README_MISSING_SUBSECTION, EXPECTED_ERROR_README_MISSING_SUBSECTION),
        (README_MISSING_TEXT, EXPECTED_ERROR_README_MISSING_TEXT),
        (README_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_WRONG_FIRST_LEVEL),
        (README_MULTIPLE_WRONG_FIRST_LEVEL, EXPECTED_ERROR_README_MULTIPLE_WRONG_FIRST_LEVEL),
        (README_MISSING_CONTENT, EXPECTED_ERROR_README_MISSING_CONTENT),
    ]				,   )
def 					__lowerCAmelCase		(  __lowerCamelCase	:    List[str]				,   __lowerCamelCase	:    int	)  ->						List[Any]:





					with tempfile.TemporaryDirectory() as tmp_dir:
										__lowerCAmelCase	   =Path(A_	) / """README.md"""
										with open(A_				,   """w+"""	) as readme_file:
															readme_file.write(A_	)
										__lowerCAmelCase	   =expected_error.format(path=A_	)
										with pytest.raises(A_				,   match=re.escape(A_	)	):
															__lowerCAmelCase	   =ReadMe.from_readme(A_				,   A_	)
															readme.validate()






@pytest.mark.parametrize(
    """readme_md, expected_error"""				,   [
        (README_MULTIPLE_SAME_HEADING_1, EXPECTED_ERROR_README_MULTIPLE_SAME_HEADING_1),
    ]				,   )
def 					__lowerCAmelCase		(  __lowerCamelCase	:    Any				,   __lowerCamelCase	:    Union[str, Any]	)  ->						Union[str, Any]:





					with tempfile.TemporaryDirectory() as tmp_dir:
										__lowerCAmelCase	   =Path(A_	) / """README.md"""
										with open(A_				,   """w+"""	) as readme_file:
															readme_file.write(A_	)
										__lowerCAmelCase	   =expected_error.format(path=A_	)
										with pytest.raises(A_				,   match=re.escape(A_	)	):
															ReadMe.from_readme(A_				,   A_	)






@pytest.mark.parametrize(
    """readme_md,"""				,   [
        (README_MULTIPLE_SAME_HEADING_1),
    ]				,   )
def 					__lowerCAmelCase		(  __lowerCamelCase	:    int	)  ->						Optional[Any]:





					with tempfile.TemporaryDirectory() as tmp_dir:
										__lowerCAmelCase	   =Path(A_	) / """README.md"""
										with open(A_				,   """w+"""	) as readme_file:
															readme_file.write(A_	)
										ReadMe.from_readme(A_				,   A_				,   suppress_parsing_errors=A_	)









import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import numpy as np
import pytest

from datasets.arrow_dataset import Dataset
from datasets.search import ElasticSearchIndex, FaissIndex, MissingIndex

from .utils import require_elasticsearch, require_faiss


lowercase_	       =			pytest.mark.integration






@require_faiss
class 						__a (     SCREAMING_SNAKE_CASE				):



	def 		UpperCamelCase							(						self							:    Optional[int])->							List[Any]:
						__lowerCAmelCase	   =Dataset.from_dict({"""filename""": ["""my_name-train""" + """_""" + str(snake_case_) for x in np.arange(30).tolist()]})
						return dset



	def 		UpperCamelCase							(						self							:    Optional[Any])->							str:
						import faiss

						__lowerCAmelCase	   =self._create_dummy_dataset()
						__lowerCAmelCase	   =dset.map(
						    lambda snake_case_				,	snake_case_: {"vecs": i * np.ones(5				,	dtype=np.floataa)}				,	with_indices=snake_case_				,	keep_in_memory=snake_case_)
						__lowerCAmelCase	   =dset.add_faiss_index("""vecs"""				,	batch_size=1_00				,	metric_type=faiss.METRIC_INNER_PRODUCT)
						__lowerCAmelCase   ,	__lowerCAmelCase	   =dset.get_nearest_examples("""vecs"""				,	np.ones(5				,	dtype=np.floataa))
						self.assertEqual(examples["""filename"""][0]				,	"""my_name-train_29""")
						dset.drop_index("""vecs""")



	def 		UpperCamelCase							(						self							:    Tuple)->							int:
						import faiss

						__lowerCAmelCase	   =self._create_dummy_dataset()
						dset.add_faiss_index_from_external_arrays(
						    external_arrays=np.ones((30, 5)) * np.arange(30).reshape(-1				,	1)				,	index_name="""vecs"""				,	batch_size=1_00				,	metric_type=faiss.METRIC_INNER_PRODUCT				,	)
						__lowerCAmelCase   ,	__lowerCAmelCase	   =dset.get_nearest_examples("""vecs"""				,	np.ones(5				,	dtype=np.floataa))
						self.assertEqual(examples["""filename"""][0]				,	"""my_name-train_29""")



	def 		UpperCamelCase							(						self							:    int)->							List[str]:
						import faiss

						__lowerCAmelCase	   =self._create_dummy_dataset()
						dset.add_faiss_index_from_external_arrays(
						    external_arrays=np.ones((30, 5)) * np.arange(30).reshape(-1				,	1)				,	index_name="""vecs"""				,	metric_type=faiss.METRIC_INNER_PRODUCT				,	)

						# Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
						# ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
						# see https://bugs.python.org/issue14243 and
						# https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
						with tempfile.NamedTemporaryFile(delete=snake_case_) as tmp_file:
											dset.save_faiss_index("""vecs"""				,	tmp_file.name)
											dset.load_faiss_index("""vecs2"""				,	tmp_file.name)
						os.unlink(tmp_file.name)

						__lowerCAmelCase   ,	__lowerCAmelCase	   =dset.get_nearest_examples("""vecs2"""				,	np.ones(5				,	dtype=np.floataa))
						self.assertEqual(examples["""filename"""][0]				,	"""my_name-train_29""")



	def 		UpperCamelCase							(						self							:    Any)->							Optional[int]:
						__lowerCAmelCase	   =self._create_dummy_dataset()
						dset.add_faiss_index_from_external_arrays(
						    external_arrays=np.ones((30, 5)) * np.arange(30).reshape(-1				,	1)				,	index_name="""vecs""")
						dset.drop_index("""vecs""")
						self.assertRaises(snake_case_				,	partial(dset.get_nearest_examples				,	"""vecs2"""				,	np.ones(5				,	dtype=np.floataa)))



	def 		UpperCamelCase							(						self							:    Optional[int])->							List[str]:
						from elasticsearch import Elasticsearch

						__lowerCAmelCase	   =self._create_dummy_dataset()
						with patch("""elasticsearch.Elasticsearch.search""") as mocked_search, patch(
						    """elasticsearch.client.IndicesClient.create""") as mocked_index_create, patch("""elasticsearch.helpers.streaming_bulk""") as mocked_bulk:
											__lowerCAmelCase	   ={"""acknowledged""": True}
											mocked_bulk.return_value([(True, None)] * 30)
											__lowerCAmelCase	   ={"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 29}]}}
											__lowerCAmelCase	   =Elasticsearch()

											dset.add_elasticsearch_index("""filename"""				,	es_client=snake_case_)
											__lowerCAmelCase   ,	__lowerCAmelCase	   =dset.get_nearest_examples("""filename"""				,	"""my_name-train_29""")
											self.assertEqual(examples["""filename"""][0]				,	"""my_name-train_29""")






@require_faiss
class 						__a (     SCREAMING_SNAKE_CASE				):



	def 		UpperCamelCase							(						self							:    Dict)->							int:
						import faiss

						__lowerCAmelCase	   =FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT)

						# add vectors
						index.add_vectors(np.eye(5				,	dtype=np.floataa))
						self.assertIsNotNone(index.faiss_index)
						self.assertEqual(index.faiss_index.ntotal				,	5)
						index.add_vectors(np.zeros((5, 5)				,	dtype=np.floataa))
						self.assertEqual(index.faiss_index.ntotal				,	10)

						# single query
						__lowerCAmelCase	   =np.zeros(5				,	dtype=np.floataa)
						__lowerCAmelCase	   =1
						__lowerCAmelCase   ,	__lowerCAmelCase	   =index.search(snake_case_)
						self.assertRaises(snake_case_				,	index.search				,	query.reshape(-1				,	1))
						self.assertGreater(scores[0]				,	0)
						self.assertEqual(indices[0]				,	1)

						# batched queries
						__lowerCAmelCase	   =np.eye(5				,	dtype=np.floataa)[::-1]
						__lowerCAmelCase   ,	__lowerCAmelCase	   =index.search_batch(snake_case_)
						self.assertRaises(snake_case_				,	index.search_batch				,	queries[0])
						__lowerCAmelCase	   =[scores[0] for scores in total_scores]
						__lowerCAmelCase	   =[indices[0] for indices in total_indices]
						self.assertGreater(np.min(snake_case_)				,	0)
						self.assertListEqual([4, 3, 2, 1, 0]				,	snake_case_)



	def 		UpperCamelCase							(						self							:    Optional[Any])->							str:
						import faiss

						__lowerCAmelCase	   =FaissIndex(string_factory="""Flat""")
						index.add_vectors(np.eye(5				,	dtype=np.floataa))
						self.assertIsInstance(index.faiss_index				,	faiss.IndexFlat)
						__lowerCAmelCase	   =FaissIndex(string_factory="""LSH""")
						index.add_vectors(np.eye(5				,	dtype=np.floataa))
						self.assertIsInstance(index.faiss_index				,	faiss.IndexLSH)
						with self.assertRaises(snake_case_):
											__lowerCAmelCase	   =FaissIndex(string_factory="""Flat"""				,	custom_index=faiss.IndexFlat(5))



	def 		UpperCamelCase							(						self							:    List[str])->							Any:
						import faiss

						__lowerCAmelCase	   =faiss.IndexFlat(5)
						__lowerCAmelCase	   =FaissIndex(custom_index=snake_case_)
						index.add_vectors(np.eye(5				,	dtype=np.floataa))
						self.assertIsInstance(index.faiss_index				,	faiss.IndexFlat)



	def 		UpperCamelCase							(						self							:    str)->							Union[str, Any]:
						import faiss

						__lowerCAmelCase	   =FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT)
						index.add_vectors(np.eye(5				,	dtype=np.floataa))

						# Setting delete=False and unlinking manually is not pretty... but it is required on Windows to
						# ensure somewhat stable behaviour. If we don't, we get PermissionErrors. This is an age-old issue.
						# see https://bugs.python.org/issue14243 and
						# https://stackoverflow.com/questions/23212435/permission-denied-to-write-to-my-temporary-file/23212515
						with tempfile.NamedTemporaryFile(delete=snake_case_) as tmp_file:
											index.save(tmp_file.name)
											__lowerCAmelCase	   =FaissIndex.load(tmp_file.name)
						os.unlink(tmp_file.name)

						__lowerCAmelCase	   =np.zeros(5				,	dtype=np.floataa)
						__lowerCAmelCase	   =1
						__lowerCAmelCase   ,	__lowerCAmelCase	   =index.search(snake_case_)
						self.assertGreater(scores[0]				,	0)
						self.assertEqual(indices[0]				,	1)






@require_faiss
def 					__lowerCAmelCase		(  __lowerCamelCase	:    Optional[Any]	)  ->						List[str]:
					import faiss

					__lowerCAmelCase	   =FaissIndex(metric_type=faiss.METRIC_INNER_PRODUCT	)
					index.add_vectors(np.eye(5				,   dtype=np.floataa	)	)

					__lowerCAmelCase	   ="""index.faiss"""
					__lowerCAmelCase	   =f"""mock://{index_name}"""
					index.save(__lowerCamelCase				,   storage_options=mockfs.storage_options	)
					__lowerCAmelCase	   =FaissIndex.load(__lowerCamelCase				,   storage_options=mockfs.storage_options	)

					__lowerCAmelCase	   =np.zeros(5				,   dtype=np.floataa	)
					__lowerCAmelCase	   =1
					__lowerCAmelCase   ,	__lowerCAmelCase	   =index.search(__lowerCamelCase	)
					assert scores[0] > 0
					assert indices[0] == 1






@require_elasticsearch
class 						__a (     SCREAMING_SNAKE_CASE				):



	def 		UpperCamelCase							(						self							:    List[str])->							Optional[int]:
						from elasticsearch import Elasticsearch

						with patch("""elasticsearch.Elasticsearch.search""") as mocked_search, patch(
						    """elasticsearch.client.IndicesClient.create""") as mocked_index_create, patch("""elasticsearch.helpers.streaming_bulk""") as mocked_bulk:
											__lowerCAmelCase	   =Elasticsearch()
											__lowerCAmelCase	   ={"""acknowledged""": True}
											__lowerCAmelCase	   =ElasticSearchIndex(es_client=snake_case_)
											mocked_bulk.return_value([(True, None)] * 3)
											index.add_documents(["""foo""", """bar""", """foobar"""])

											# single query
											__lowerCAmelCase	   ="""foo"""
											__lowerCAmelCase	   ={"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 0}]}}
											__lowerCAmelCase   ,	__lowerCAmelCase	   =index.search(snake_case_)
											self.assertEqual(scores[0]				,	1)
											self.assertEqual(indices[0]				,	0)

											# single query with timeout
											__lowerCAmelCase	   ="""foo"""
											__lowerCAmelCase	   ={"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 0}]}}
											__lowerCAmelCase   ,	__lowerCAmelCase	   =index.search(snake_case_				,	request_timeout=30)
											self.assertEqual(scores[0]				,	1)
											self.assertEqual(indices[0]				,	0)

											# batched queries
											__lowerCAmelCase	   =["""foo""", """bar""", """foobar"""]
											__lowerCAmelCase	   ={"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 1}]}}
											__lowerCAmelCase   ,	__lowerCAmelCase	   =index.search_batch(snake_case_)
											__lowerCAmelCase	   =[scores[0] for scores in total_scores]
											__lowerCAmelCase	   =[indices[0] for indices in total_indices]
											self.assertGreater(np.min(snake_case_)				,	0)
											self.assertListEqual([1, 1, 1]				,	snake_case_)

											# batched queries with timeout
											__lowerCAmelCase	   =["""foo""", """bar""", """foobar"""]
											__lowerCAmelCase	   ={"""hits""": {"""hits""": [{"""_score""": 1, """_id""": 1}]}}
											__lowerCAmelCase   ,	__lowerCAmelCase	   =index.search_batch(snake_case_				,	request_timeout=30)
											__lowerCAmelCase	   =[scores[0] for scores in total_scores]
											__lowerCAmelCase	   =[indices[0] for indices in total_indices]
											self.assertGreater(np.min(snake_case_)				,	0)
											self.assertListEqual([1, 1, 1]				,	snake_case_)




