def test_put_node(self):
stub = IndexStub(self.channel)
# get node
request = GetStatusRequest()
response = stub.GetStatus(request)
self.assertEqual(
0,
pickle.loads(response.node_status)['partner_nodes_count'])
self.assertEqual(True, response.status.success)
# put node
request = PutNodeRequest()
request.node_name = 'localhost:{0}'.format(get_free_port())
response = stub.PutNode(request)
sleep(1) # wait for node to be added
self.assertEqual(True, response.status.success)
# get node
request = GetStatusRequest()
response = stub.GetStatus(request)
self.assertEqual(
1,
pickle.loads(response.node_status)['partner_nodes_count'])
self.assertEqual(True, response.status.success)
def test_get_snapshot(self):
stub = IndexStub(self.channel)
# create snapshot
request = CreateSnapshotRequest()
response = stub.CreateSnapshot(request)
sleep(1) # wait for snapshot file to be created
self.assertEqual(True, response.status.success)
self.assertEqual(True,
os.path.exists(self.indexer.get_snapshot_file_name()))
with zipfile.ZipFile(self.indexer.get_snapshot_file_name()) as f:
self.assertEqual(['raft.bin'], f.namelist())
# get snapshot
request = GetSnapshotRequest()
request.chunk_size = 1024
response = stub.GetSnapshot(request)
download_file_name = self.temp_dir.name + '/snapshot_downloaded.zip'
with open(download_file_name, 'wb') as f:
for snapshot in response:
f.write(snapshot.chunk)
with zipfile.ZipFile(download_file_name) as f:
self.assertEqual(['raft.bin'], f.namelist())
def test_optimize_index(self):
stub = IndexStub(self.channel)
# read index_config.yaml
with open(self.example_dir + '/index_config.yaml',
'r',
encoding='utf-8') as file_obj:
index_config_dict = yaml.safe_load(file_obj.read())
# create index
request = CreateIndexRequest()
request.index_name = 'test_index'
request.index_config = pickle.dumps(index_config_dict)
request.sync = True
response = stub.CreateIndex(request)
self.assertEqual(True, response.status.success)
# get index
request = GetIndexRequest()
request.index_name = 'test_index'
response = stub.GetIndex(request)
self.assertEqual(True, response.status.success)
self.assertEqual('test_index', response.index_stats.name)
# optimize index
request = OptimizeIndexRequest()
request.index_name = 'test_index'
request.sync = True
response = stub.CloseIndex(request)
self.assertEqual(True, response.status.success)
def test_is_ready(self):
stub = IndexStub(self.channel)
# is the cluster ready
request = IsReadyRequest()
response = stub.IsReady(request)
self.assertEqual(True, response.status.success)
def test_is_alive(self):
stub = IndexStub(self.channel)
# is the node alive
request = IsAliveRequest()
response = stub.IsAlive(request)
self.assertEqual(True, response.status.success)
def test_get_status(self):
stub = IndexStub(self.channel)
# get node
request = GetStatusRequest()
response = stub.GetStatus(request)
self.assertEqual(
0,
pickle.loads(response.node_status)['partner_nodes_count'])
self.assertEqual(True, response.status.success)
def test_create_snapshot(self):
stub = IndexStub(self.channel)
self.assertEqual(False,
os.path.exists(self.indexer.get_snapshot_file_name()))
# create snapshot
request = CreateSnapshotRequest()
response = stub.CreateSnapshot(request)
sleep(1) # wait for snapshot file to be created
self.assertEqual(True, response.status.success)
self.assertEqual(True,
os.path.exists(self.indexer.get_snapshot_file_name()))
with zipfile.ZipFile(self.indexer.get_snapshot_file_name()) as f:
self.assertEqual(['raft.bin'], f.namelist())
def test_get_document(self):
stub = IndexStub(self.channel)
# read index_config.yaml
with open(self.example_dir + '/index_config.yaml',
'r',
encoding='utf-8') as file_obj:
index_config_dict = yaml.safe_load(file_obj.read())
# create index
request = CreateIndexRequest()
request.index_name = 'test_index'
request.index_config = pickle.dumps(index_config_dict)
request.sync = True
response = stub.CreateIndex(request)
self.assertEqual(True, response.status.success)
# get index
request = GetIndexRequest()
request.index_name = 'test_index'
response = stub.GetIndex(request)
self.assertEqual(True, response.status.success)
self.assertEqual('test_index', response.index_stats.name)
# read doc1.yaml
with open(self.example_dir + '/doc1.yaml', 'r',
encoding='utf-8') as file_obj:
fields_dict = yaml.safe_load(file_obj.read())
# put document
request = PutDocumentRequest()
request.index_name = 'test_index'
request.doc_id = '1'
request.fields = pickle.dumps(fields_dict)
request.sync = True
response = stub.PutDocument(request)
self.assertEqual(1, response.count)
self.assertEqual(True, response.status.success)
# commit
request = CommitIndexRequest()
request.index_name = 'test_index'
request.sync = True
response = stub.CommitIndex(request)
self.assertEqual(True, response.status.success)
# get document
request = GetDocumentRequest()
request.index_name = 'test_index'
request.doc_id = '1'
response = stub.GetDocument(request)
self.assertEqual(True, response.status.success)
self.assertEqual('1', pickle.loads(response.fields)['id'])
self.assertEqual('Search engine (computing)',
pickle.loads(response.fields)['title'])
def test_snapshot_exists(self):
stub = IndexStub(self.channel)
# snapshot exists
request = IsSnapshotExistRequest()
response = stub.IsSnapshotExist(request)
self.assertEqual(True, response.status.success)
self.assertEqual(False, response.exist)
# create snapshot
request = CreateSnapshotRequest()
response = stub.CreateSnapshot(request)
sleep(1) # wait for snapshot file to be created
self.assertEqual(True, response.status.success)
self.assertEqual(True,
os.path.exists(self.indexer.get_snapshot_file_name()))
# snapshot exists
request = IsSnapshotExistRequest()
response = stub.IsSnapshotExist(request)
self.assertEqual(True, response.status.success)
self.assertEqual(True, response.exist)
def test_search_documents(self):
stub = IndexStub(self.channel)
# read index_config.yaml
with open(self.example_dir + '/index_config.yaml',
'r',
encoding='utf-8') as file_obj:
index_config_dict = yaml.safe_load(file_obj.read())
# create index
request = CreateIndexRequest()
request.index_name = 'test_index'
request.index_config = pickle.dumps(index_config_dict)
request.sync = True
response = stub.CreateIndex(request)
self.assertEqual(True, response.status.success)
# get index
request = GetIndexRequest()
request.index_name = 'test_index'
response = stub.GetIndex(request)
self.assertEqual(True, response.status.success)
self.assertEqual('test_index', response.index_stats.name)
# read bulk_put.yaml
with open(self.example_dir + '/bulk_put.yaml', 'r',
encoding='utf-8') as file_obj:
docs_dict = yaml.safe_load(file_obj.read())
# put documents
request = PutDocumentsRequest()
request.index_name = 'test_index'
request.docs = pickle.dumps(docs_dict)
request.sync = True
response = stub.PutDocuments(request)
self.assertEqual(5, response.count)
self.assertEqual(True, response.status.success)
# commit
request = CommitIndexRequest()
request.index_name = 'test_index'
request.sync = True
response = stub.CommitIndex(request)
self.assertEqual(True, response.status.success)
# read weighting.yaml
with open(self.example_dir + '/weighting.yaml', 'r',
encoding='utf-8') as file_obj:
weighting_dict = yaml.safe_load(file_obj.read())
# search documents
request = SearchDocumentsRequest()
request.index_name = 'test_index'
request.query = 'search'
request.search_field = 'text'
request.page_num = 1
request.page_len = 10
request.weighting = pickle.dumps(weighting_dict)
response = stub.SearchDocuments(request)
self.assertEqual(5, pickle.loads(response.results)['total'])
self.assertEqual(True, response.status.success)