-
Notifications
You must be signed in to change notification settings - Fork 41
/
Copy pathtest_write_insert.py
130 lines (121 loc) · 5.75 KB
/
test_write_insert.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
"""Test the write database.
"""
import json
import pytest
import os
import time
import numpy as np
from vectorai.models.deployed import ViText2Vec
from vectorai.write import ViWriteClient
from vectorai.errors import APIError, MissingFieldError, MissingFieldWarning, CollectionNameError
from vectorai.client import ViClient
from .utils import TempClientWithDocs
class TestInsert:
@pytest.mark.use_client
def test_insert_documents_simple_and_collection_stats_match(self, test_client,
test_collection_name):
"""
Testing for simple document insertion
"""
if test_collection_name in test_client.list_collections():
test_client.delete_collection(test_collection_name)
sample_documents = test_client.create_sample_documents(10)
test_client.insert_documents(test_collection_name, sample_documents)
time.sleep(10)
assert test_client.collection_stats(test_collection_name)['number_of_documents'] == 10
test_client.delete_collection(test_collection_name)
time.sleep(3)
@pytest.mark.use_client
def test_inserting_documents_without_id_fields(self, test_client, test_collection_name):
"""
Test inserting documents if they do not have an ID field.
"""
if test_collection_name in test_client.list_collections():
test_client.delete_collection(test_collection_name)
sample_documents = test_client.create_sample_documents(10)
# Remove the ID fields
{x.pop('_id') for x in sample_documents}
test_client.insert_documents(test_collection_name, sample_documents)
time.sleep(10)
assert test_client.collection_stats(test_collection_name)['number_of_documents'] == 10
test_client.delete_collection(test_collection_name)
time.sleep(3)
@pytest.mark.use_client
def test_inserting_documents_without_id_fields_with_overwrite(self, test_client,
test_collection_name):
"""
Test inserting documents if they do not have an ID field.
"""
if test_collection_name in test_client.list_collections():
test_client.delete_collection(test_collection_name)
sample_documents = test_client.create_sample_documents(10)
# Remove the ID fields
{x.pop('_id') for x in sample_documents}
with pytest.warns(MissingFieldWarning):
test_client.insert_documents(test_collection_name, sample_documents, overwrite=True)
time.sleep(10)
assert test_client.collection_stats(test_collection_name)['number_of_documents'] == 10
test_client.delete_collection(test_collection_name)
time.sleep(3)
@pytest.mark.use_client
def test_inserting_documents_when_id_is_not_a_string(self, test_client, test_collection_name):
"""
Test inserting documents when ID is not a string
"""
if test_collection_name in test_client.list_collections():
test_client.delete_collection(test_collection_name)
sample_documents = test_client.create_sample_documents(10)
# Create integer IDs strings
{x.update({'_id': int(x['_id'])}) for x in sample_documents}
test_client.insert_documents(test_collection_name, sample_documents, overwrite=False)
time.sleep(10)
assert test_client.collection_stats(test_collection_name)['number_of_documents'] == 10
test_client.delete_collection(test_collection_name)
time.sleep(3)
@pytest.mark.use_client
def test_inserting_documents_when_id_is_not_a_string_with_overwrite(self, test_client,
test_collection_name):
"""
Test inserting documents when ID is not a string
"""
if test_collection_name in test_client.list_collections():
test_client.delete_collection(test_collection_name)
sample_documents = test_client.create_sample_documents(10)
# Create integer IDs strings
{x.update({'_id': int(x['_id'])}) for x in sample_documents}
test_client.insert_documents(test_collection_name, sample_documents, overwrite=True)
time.sleep(10)
assert test_client.collection_stats(test_collection_name)['number_of_documents'] == 10
test_client.delete_collection(test_collection_name)
time.sleep(3)
@pytest.mark.use_client
def test_insert_single_document(self, test_client, test_collection_name):
if test_collection_name not in test_client.list_collections():
test_client.create_collection(test_collection_name)
# document = {"sample_vector_": test_client.generate_vector(20), "sample_name": "hi"}
document = test_client.create_sample_document(1)
response = test_client.insert(
collection_name=test_collection_name, document=document
)
assert response['status'] == 'success'
@pytest.mark.use_client
def test_insert_single_document_error(self, test_client, test_collection_name):
"""Trigger an insert fail error
"""
with pytest.raises(APIError):
if test_collection_name not in test_client.list_collections():
test_client.create_collection(test_collection_name)
document = {
"sample_vectors_": [test_client.generate_vector(20)] + [np.nan],
"samplename": [["hi"]],
}
response = test_client.insert(
collection_name=test_collection_name, document=document
)
@pytest.mark.use_client
def test_clean_up(self, test_client, test_collection_name):
"""Remove a collection if it is there.
"""
if test_collection_name in test_client.list_collections():
test_client.delete_collection(test_collection_name)
assert test_collection_name not in test_client.list_collections()