11import pytest
2- from testcontainers .core .container import DockerContainer # type: ignore
3- from testcontainers .core .waiting_utils import wait_for_logs # type: ignore
2+ from testcontainers .qdrant import QdrantContainer
43from qdrant_client import QdrantClient , models
54import uuid
65from pyspark .sql import SparkSession
76from typing import NamedTuple
7+ from uuid import uuid4
88
99
1010QDRANT_GRPC_PORT = 6334
1111QDRANT_EMBEDDING_DIM = 6
1212QDRANT_DISTANCE = models .Distance .COSINE
13+ QDRANT_API_KEY = uuid4 ().hex
1314
1415
1516class Qdrant (NamedTuple ):
1617 url : str
18+ api_key : str
1719 collection_name : str
1820 client : QdrantClient
1921
2022
21- qdrant_container = DockerContainer ( "qdrant/qdrant" ). with_exposed_ports ( QDRANT_GRPC_PORT )
23+ qdrant_container = QdrantContainer ( image = "qdrant/qdrant:latest" , api_key = QDRANT_API_KEY )
2224
2325
2426# Reference: https://gist.github.com/dizzythinks/f3bb37fd8ab1484bfec79d39ad8a92d3
@@ -36,9 +38,6 @@ def get_pom_version():
3638@pytest .fixture (scope = "module" , autouse = True )
3739def setup_container (request ):
3840 qdrant_container .start ()
39- wait_for_logs (
40- qdrant_container , ".*Actix runtime found; starting in Actix runtime.*" , 60
41- )
4241
4342 def remove_container ():
4443 qdrant_container .stop ()
@@ -70,6 +69,8 @@ def qdrant():
7069 host = host ,
7170 grpc_port = grpc_port ,
7271 prefer_grpc = True ,
72+ api_key = QDRANT_API_KEY ,
73+ https = False ,
7374 )
7475
7576 collection_name = str (uuid .uuid4 ())
@@ -99,6 +100,7 @@ def qdrant():
99100 url = f"http://{ host } :{ grpc_port } " ,
100101 client = client ,
101102 collection_name = collection_name ,
103+ api_key = QDRANT_API_KEY ,
102104 )
103105
104106 return client .close ()
0 commit comments