tests: add VectorEntity objectbox#24

Author: loryruta

tests/common.py

@@ -1,9 +1,12 @@
 import objectbox
 import os
+from os import path
 import shutil
 import pytest
-from tests.model import TestEntity, TestEntityDatetime, TestEntityFlex
 import numpy as np
+from typing import *
+from tests.model import *
+
 
 test_dir = 'testdata'
 
@@ -23,15 +26,13 @@ def autocleanup():
 
 def load_empty_test_objectbox(db_name: str = test_dir) -> objectbox.ObjectBox:
     model = objectbox.Model()
-    from objectbox.model import IdUid
     model.entity(TestEntity, last_property_id=IdUid(27, 1027))
     model.last_entity_id = IdUid(2, 2)
 
     return objectbox.Builder().model(model).directory(db_name).build()
 
 def load_empty_test_datetime(name: str = "") -> objectbox.ObjectBox:
     model = objectbox.Model()
-    from objectbox.model import IdUid
     model.entity(TestEntityDatetime, last_property_id=IdUid(4, 2004))
     model.last_entity_id = IdUid(2, 2)
 
@@ -42,7 +43,6 @@ def load_empty_test_datetime(name: str = "") -> objectbox.ObjectBox:
 
 def load_empty_test_flex(name: str = "") -> objectbox.ObjectBox:
     model = objectbox.Model()
-    from objectbox.model import IdUid
     model.entity(TestEntityFlex, last_property_id=IdUid(3, 3003))
     model.last_entity_id = IdUid(3, 3)
 
@@ -51,6 +51,26 @@ def load_empty_test_flex(name: str = "") -> objectbox.ObjectBox:
     return objectbox.Builder().model(model).directory(db_name).build()
 
 
+def create_test_objectbox(db_name: Optional[str] = None, clear_db: bool = True) -> objectbox.ObjectBox:
+    """ Creates an ObjectBox instance with all entities. """
+
+    db_path = test_dir if db_name is None else path.join(test_dir, db_name)
+    print(f"DB path: \"{db_path}\"")
+
+    if clear_db and path.exists(db_path):
+        shutil.rmtree(db_path)
+
+    model = objectbox.Model()
+    model.entity(TestEntity, last_property_id=IdUid(27, 1027))
+    model.entity(TestEntityDatetime, last_property_id=IdUid(4, 2004))
+    model.entity(TestEntityFlex, last_property_id=IdUid(3, 3003))
+    model.entity(VectorEntity, last_property_id=IdUid(3, 4003))
+    model.last_entity_id = IdUid(4, 4)
+    model.last_index_id = IdUid(3, 40001)
+
+    return objectbox.Builder().model(model).directory(db_path).build()
+
+
 def assert_equal_prop(actual, expected, default):
     assert actual == expected or (isinstance(
         expected, objectbox.model.Property) and actual == default)

tests/model.py

@@ -52,11 +52,23 @@ class TestEntityDatetime:
     def __init__(self, string: str = ""):
         self.str = string
 
+
 @Entity(id=3, uid=3)
 class TestEntityFlex:
     id = Id(id=1, uid=3001)
     flex_dict = Property(Dict[str, Any], type=PropertyType.flex, id=2, uid=3002)
     flex_int = Property(int, type=PropertyType.flex, id=3, uid=3003)
 
     def __init__(self, string: str = ""):
-        self.str = string
+        self.str = string
+
+
+@Entity(id=4, uid=4)
+class VectorEntity:
+    id = Id(id=1, uid=4001)
+    name = Property(str, type=PropertyType.string, id=2, uid=4002)
+    vector = Property(np.ndarray, type=PropertyType.floatVector, id=3, uid=4003,
+                      index=HnswIndex(
+                          id=3, uid=40001,
+                          dimensions=2, distance_type=HnswDistanceType.EUCLIDEAN)
+                      )