test: refactor querylang test use types (#1963)

jina-ai · Feb 18, 2021 · b1fab67 · b1fab67
1 parent c379b09
commit b1fab67
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 16 deletions.
diff --git a/tests/unit/drivers/querylang/test_querylang_drivers.py b/tests/unit/drivers/querylang/test_querylang_drivers.py
@@ -2,41 +2,40 @@
 
 from jina.executors.segmenters import BaseSegmenter
 from jina.flow import Flow
-from jina.proto import jina_pb2
+from jina import Document
 
 
 def random_docs(num_docs):
     for j in range(num_docs):
-        d = jina_pb2.DocumentProto()
+        d = Document()
         d.text = 'hello world'
-        d.uri = 'doc://'
         d.tags['id'] = j
         for m in range(10):
-            dm = d.matches.add()
+            dm = Document()
             dm.text = 'match to other world'
-            dm.uri = 'doc://match'
             dm.tags['id'] = m
+            d.matches.add(dm)
         yield d
 
 
 def random_docs_to_chunk():
-    d1 = jina_pb2.DocumentProto()
+    d1 = Document()
     d1.tags['id'] = 1
     d1.text = 'chunk1 chunk2'
     yield d1
-    d2 = jina_pb2.DocumentProto()
+    d2 = Document()
     d2.tags['id'] = 1
     d2.text = 'chunk3'
     yield d2
 
 
 def random_docs_with_tags():
-    d1 = jina_pb2.DocumentProto()
+    d1 = Document()
     d1.tags['id'] = 1
     d1.text = 'a'
     d1.tags.update({'id': 1})
     yield d1
-    d2 = jina_pb2.DocumentProto()
+    d2 = Document()
     d2.tags['id'] = 2
     d2.tags.update({'id': 2})
     d2.text = 'b'

diff --git a/tests/unit/drivers/querylang/test_querylang_reader.py b/tests/unit/drivers/querylang/test_querylang_reader.py
@@ -6,29 +6,28 @@
 from jina.drivers.querylang.slice import SliceQL
 from jina.drivers.querylang.select import ExcludeQL
 from jina.flow import Flow
-from jina.proto import jina_pb2
+from jina import Document
 from jina.types.querylang import QueryLang
 from jina.types.sets import QueryLangSet
 
 
 def random_docs(num_docs):
     for j in range(num_docs):
-        d = jina_pb2.DocumentProto()
+        d = Document()
         d.tags['id'] = j
         d.text = 'hello world'
-        d.uri = 'doc://'
         for m in range(10):
-            dm = d.matches.add()
+            dm = Document()
             dm.text = 'match to hello world'
-            dm.uri = 'doc://match'
             dm.tags['id'] = m
             dm.score.ref_id = d.id
+            d.matches.add(dm)
             for mm in range(10):
-                dmm = dm.matches.add()
+                dmm = Document()
                 dmm.text = 'nested match to match'
-                dmm.uri = 'doc://match/match'
                 dmm.tags['id'] = mm
                 dmm.score.ref_id = dm.id
+                dm.matches.add(dmm)
         yield d