fixed sample tests

Author: daniel-sanche

datastore/cloud-client/vector_search.py

@@ -30,7 +30,7 @@ def store_vectors():
 
     client.put(entity)
     # [END datastore_store_vectors]
-    return client
+    return client, entity
 
 
 def vector_search_basic(db):

datastore/cloud-client/vector_search_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import pytest
 
 from google.cloud import datastore
 from google.cloud.datastore.vector import Vector
@@ -28,12 +29,21 @@
 PROJECT_ID = os.environ["GOOGLE_CLOUD_PROJECT"]
 
 
-def test_store_vectors():
-    client = store_vectors()
-
-    results = client.query("coffee-beans", limit=5).fetch()
+@pytest.fixture(scope="module")
+def db():
+    client = datastore.Client()
+    _clear_db(client)
+    entity_list = add_coffee_beans_data(client)
+    yield client
+    for e in entity_list:
+        client.delete(e)
 
-    assert len(list(results)) == 1
+def _clear_db(db):
+    """remove all entities with kind-coffee-beans, so we have a new databse"""
+    query = db.query(kind="coffee-beans")
+    query.keys_only()
+    keys = list(query.fetch())
+    db.delete_multi(keys)
 
 
 def add_coffee_beans_data(db):
@@ -46,60 +56,54 @@ def add_coffee_beans_data(db):
     entity4 = datastore.Entity(db.key("coffee-beans", "Liberica"))
     entity4.update({"embedding_field": Vector([3.0, 1.0, 2.0]), "color": "green"})
 
-    db.put_multi([entity1, entity2, entity3, entity4])
+    entity_list = [entity1, entity2, entity3, entity4]
+    db.put_multi(entity_list)
+    return entity_list
 
+def test_store_vectors():
+    # run an ensure there are no exceptions
+    client, entity = store_vectors()
+    client.delete(entity)
 
-def test_vector_search_basic():
-    db = datastore.Client()
-    add_coffee_beans_data(db)
-
+def test_vector_search_basic(db):
     vector_query = vector_search_basic(db)
     results = list(vector_query.fetch())
 
     assert len(results) == 4
-    assert results[0].name == "Liberica"
-    assert results[1].name == "Robusta"
-    assert results[2].name == "Arabica"
-    assert results[3].name == "Excelsa"
+    assert results[0].key.name == "Liberica"
+    assert results[1].key.name == "Robusta"
+    assert results[2].key.name == "Arabica"
+    assert results[3].key.name == "Excelsa"
 
 
-def test_vector_search_prefilter():
-    db = datastore.Client()
-    add_coffee_beans_data(db)
-
+def test_vector_search_prefilter(db):
     vector_query = vector_search_prefilter(db)
     results = list(vector_query.fetch())
 
     assert len(results) == 2
-    assert results[0].name == "Arabica"
-    assert results[1].name == "Excelsa"
-
+    assert results[0].key.name == "Arabica"
+    assert results[1].key.name == "Excelsa"
 
-def test_vector_search_distance_result_field():
-    db = datastore.Client()
-    add_coffee_beans_data(db)
 
+def test_vector_search_distance_result_field(db):
     vector_query = vector_search_distance_result_field(db)
     results = list(vector_query.fetch())
 
     assert len(results) == 4
-    assert results[0].name == "Liberica"
+    assert results[0].key.name == "Liberica"
     assert results[0]["vector_distance"] == 0.0
-    assert results[1].name == "Robusta"
+    assert results[1].key.name == "Robusta"
     assert results[1]["vector_distance"] == 1.0
-    assert results[2].name == "Arabica"
+    assert results[2].key.name == "Arabica"
     assert results[2]["vector_distance"] == 7.0
-    assert results[3].name == "Excelsa"
+    assert results[3].key.name == "Excelsa"
     assert results[3]["vector_distance"] == 8.0
 
 
-def test_vector_search_distance_threshold():
-    db = datastore.Client()
-    add_coffee_beans_data(db)
-
+def test_vector_search_distance_threshold(db):
     vector_query = vector_search_distance_threshold(db)
     results = list(vector_query.fetch())
 
     assert len(results) == 2
-    assert results[0].name == "Liberica"
-    assert results[1].name == "Robusta"
+    assert results[0].key.name == "Liberica"
+    assert results[1].key.name == "Robusta"