github.com/weaviate/weaviate@v1.24.6/test/acceptance_with_python/conftest.py (about)

     1  import os
     2  from typing import Any, Optional, List, Generator, Protocol, Type, Dict, Tuple, Union
     3  
     4  import pytest
     5  from _pytest.fixtures import SubRequest
     6  
     7  import weaviate
     8  from weaviate.collections import Collection
     9  from weaviate.collections.classes.config import (
    10      Property,
    11      _VectorizerConfigCreate,
    12      _InvertedIndexConfigCreate,
    13      _ReferencePropertyBase,
    14      _GenerativeConfigCreate,
    15      _ReplicationConfigCreate,
    16      _MultiTenancyConfigCreate,
    17      _VectorIndexConfigCreate,
    18      _RerankerConfigCreate,
    19  )
    20  from weaviate.collections.classes.types import Properties
    21  from weaviate.config import AdditionalConfig
    22  
    23  from weaviate.collections.classes.config_named_vectors import _NamedVectorConfigCreate
    24  
    25  
    26  class CollectionFactory(Protocol):
    27      """Typing for fixture."""
    28  
    29      def __call__(
    30          self,
    31          name: str = "",
    32          properties: Optional[List[Property]] = None,
    33          references: Optional[List[_ReferencePropertyBase]] = None,
    34          vectorizer_config: Optional[
    35              Union[_VectorizerConfigCreate, List[_NamedVectorConfigCreate]]
    36          ] = None,
    37          inverted_index_config: Optional[_InvertedIndexConfigCreate] = None,
    38          multi_tenancy_config: Optional[_MultiTenancyConfigCreate] = None,
    39          generative_config: Optional[_GenerativeConfigCreate] = None,
    40          headers: Optional[Dict[str, str]] = None,
    41          ports: Tuple[int, int] = (8080, 50051),
    42          data_model_properties: Optional[Type[Properties]] = None,
    43          data_model_refs: Optional[Type[Properties]] = None,
    44          replication_config: Optional[_ReplicationConfigCreate] = None,
    45          vector_index_config: Optional[_VectorIndexConfigCreate] = None,
    46          description: Optional[str] = None,
    47          reranker_config: Optional[_RerankerConfigCreate] = None,
    48      ) -> Collection[Any, Any]:
    49          """Typing for fixture."""
    50          ...
    51  
    52  
    53  @pytest.fixture
    54  def collection_factory(request: SubRequest) -> Generator[CollectionFactory, None, None]:
    55      name_fixture: Optional[str] = None
    56      client_fixture: Optional[weaviate.WeaviateClient] = None
    57  
    58      def _factory(
    59          name: str = "",
    60          properties: Optional[List[Property]] = None,
    61          references: Optional[List[_ReferencePropertyBase]] = None,
    62          vectorizer_config: Optional[
    63              Union[_VectorizerConfigCreate, List[_NamedVectorConfigCreate]]
    64          ] = None,
    65          inverted_index_config: Optional[_InvertedIndexConfigCreate] = None,
    66          multi_tenancy_config: Optional[_MultiTenancyConfigCreate] = None,
    67          generative_config: Optional[_GenerativeConfigCreate] = None,
    68          headers: Optional[Dict[str, str]] = None,
    69          ports: Tuple[int, int] = (8080, 50051),
    70          data_model_properties: Optional[Type[Properties]] = None,
    71          data_model_refs: Optional[Type[Properties]] = None,
    72          replication_config: Optional[_ReplicationConfigCreate] = None,
    73          vector_index_config: Optional[_VectorIndexConfigCreate] = None,
    74          description: Optional[str] = None,
    75          reranker_config: Optional[_RerankerConfigCreate] = None,
    76      ) -> Collection[Any, Any]:
    77          nonlocal client_fixture, name_fixture
    78          name_fixture = _sanitize_collection_name(request.node.name) + name
    79          client_fixture = weaviate.connect_to_local(
    80              headers=headers,
    81              grpc_port=ports[1],
    82              port=ports[0],
    83              additional_config=AdditionalConfig(timeout=(60, 120)),  # for image tests
    84          )
    85          client_fixture.collections.delete(name_fixture)
    86  
    87          collection: Collection[Any, Any] = client_fixture.collections.create(
    88              name=name_fixture,
    89              description=description,
    90              vectorizer_config=vectorizer_config,
    91              properties=properties,
    92              references=references,
    93              inverted_index_config=inverted_index_config,
    94              multi_tenancy_config=multi_tenancy_config,
    95              generative_config=generative_config,
    96              data_model_properties=data_model_properties,
    97              data_model_references=data_model_refs,
    98              replication_config=replication_config,
    99              vector_index_config=vector_index_config,
   100              reranker_config=reranker_config,
   101          )
   102          return collection
   103  
   104      try:
   105          yield _factory
   106      finally:
   107          if client_fixture is not None and name_fixture is not None:
   108              client_fixture.collections.delete(name_fixture)
   109              client_fixture.close()
   110  
   111  
   112  def _sanitize_collection_name(name: str) -> str:
   113      name = name.replace("[", "").replace("]", "").replace("-", "").replace(" ", "").replace(".", "")
   114      return name[0].upper() + name[1:]