fix error

gnes-ai · Aug 27, 2019 · 10f4bed · 10f4bed
1 parent 27d3c30
commit 10f4bed
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/gnes/indexer/fulltext/filesys.py b/gnes/indexer/fulltext/filesys.py
@@ -48,7 +48,7 @@ def add(self, keys: List[int], docs: List['gnes_pb2.Document'], *args, **kwargs)
                 os.makedirs(dirs)
             # keep doc meta in .meta file
             with open(os.path.join(dirs, '.meta'), 'wb') as f:
-                f.write(d.meta_info)
+                f.write(d.meta_info or b'')
 
             for i, chunk in enumerate(d.chunks):
                 with open(os.path.join(dirs, '%d.%s' % (i, self.file_suffix)), 'wb') as f:
@@ -64,13 +64,12 @@ def query(self, keys: List[int], *args, **kwargs) -> List['gnes_pb2.Document']:
             doc = gnes_pb2.Document()
             target_dirs = os.path.join(self.data_path, str(k))
 
-            with open(os.path.join(target_dirs, '.meta'), 'rb') as f:
-                doc.meta_info = f.read()
-
             if not os.path.exists(target_dirs):
                 if self.keep_na_doc:
                     res.append(self._NOT_FOUND)
             else:
+                with open(os.path.join(target_dirs, '.meta'), 'rb') as f:
+                    doc.meta_info = f.read()
                 for raw_file in os.listdir(target_dirs):
                     if not os.path.isdir(raw_file):
                         c = doc.chunks.add()

diff --git a/tests/test_dict_indexer.py b/tests/test_dict_indexer.py
@@ -20,6 +20,8 @@ def setUp(self):
         self.data_path = './test_chunkleveldb'
         self.dump_path = os.path.join(self.dirname, 'indexer.bin')
 
+        self.init_db()
+
     def tearDown(self):
         if os.path.exists(self.data_path):
             rmtree(self.data_path)
@@ -37,12 +39,12 @@ def init_db(self):
         self.db.add(list(range(len(self.video_bytes))), [self.d])
 
     def test_add_docs(self):
-        self.init_db()
+        # self.init_db()
         self.assertTrue(os.path.exists(os.path.join(self.data_path, str(self.d.doc_id))))
-        self.assertEqual(len(self.d.chunks), len(os.listdir(os.path.join(self.data_path, str(self.d.doc_id)))))
+        self.assertEqual(len(self.d.chunks), len(os.listdir(os.path.join(self.data_path, str(self.d.doc_id)))) - 1)
 
     def test_query_docs(self):
-        self.init_db()
+        # self.init_db()
 
         query_list = [0, 1, 2]
         res = self.db.query(query_list)