[Refactor]Update yatai e2e tests to the latest python APIs (#1483)

* update yatai e2e tests to the latest apis

* remove unnecessary retrieve

e2e_tests/yatai_server/test_postgres_gcs.py

@@ -1,10 +1,9 @@
 import logging
 
 from bentoml.yatai.proto.repository_pb2 import BentoUri
-from e2e_tests.cli_operations import delete_bento
+from bentoml.yatai.client import get_yatai_client
 from e2e_tests.sample_bento_service import SampleBentoService
 from e2e_tests.yatai_server.utils import (
-    get_bento_service_info,
     execute_bentoml_run_command,
     local_yatai_server,
 )
@@ -18,15 +17,16 @@ def test_yatai_server_with_postgres_and_gcs(postgres_db_container_url):
 
     with local_yatai_server(
         db_url=postgres_db_container_url, repo_base_url=gcs_bucket_name
-    ):
+    ) as yatai_service_url:
+        yc = get_yatai_client(yatai_service_url)
         logger.info('Saving bento service')
         svc = SampleBentoService()
-        svc.save()
+        svc.save(yatai_url=yatai_service_url)
         bento_tag = f'{svc.name}:{svc.version}'
         logger.info('BentoService saved')
 
         logger.info("Display bentoservice info")
-        bento = get_bento_service_info(svc.name, svc.version)
+        bento = yc.repository.get(bento_tag)
         logger.info(bento)
         assert (
             bento.uri.type == BentoUri.GCS
@@ -37,5 +37,4 @@ def test_yatai_server_with_postgres_and_gcs(postgres_db_container_url):
         assert 'cat' in run_result, 'Unexpected BentoService prediction result'
 
         logger.info(f'Deleting saved bundle {bento_tag}')
-        delete_svc_result = delete_bento(bento_tag)
-        assert f"{bento_tag} deleted" in delete_svc_result
+        yc.repository.delete(bento_tag=bento_tag)

e2e_tests/yatai_server/test_postgres_local_fs.py

@@ -1,10 +1,9 @@
 import logging
 
 from bentoml.yatai.proto.repository_pb2 import BentoUri
+from bentoml.yatai.client import get_yatai_client
 from e2e_tests.sample_bento_service import SampleBentoService
-from e2e_tests.cli_operations import delete_bento
 from e2e_tests.yatai_server.utils import (
-    get_bento_service_info,
     execute_bentoml_run_command,
     local_yatai_server,
 )
@@ -13,24 +12,26 @@
 
 
 def test_yatai_server_with_postgres_and_local_storage(postgres_db_container_url):
-    with local_yatai_server(postgres_db_container_url):
+    with local_yatai_server(postgres_db_container_url) as yatai_service_url:
+        yc = get_yatai_client(yatai_service_url)
         logger.info('Saving bento service')
         svc = SampleBentoService()
-        svc.save()
+        svc.save(yatai_url=yatai_service_url)
         bento_tag = f'{svc.name}:{svc.version}'
         logger.info('BentoService saved')
 
         logger.info("Display bentoservice info")
-        bento = get_bento_service_info(svc.name, svc.version)
+        bento = yc.repository.get(bento_tag)
         logger.info(bento)
         assert (
             bento.uri.type == BentoUri.LOCAL
         ), 'BentoService storage type mismatched, expect LOCAL'
 
         logger.info('Validate BentoService prediction result')
-        run_result = execute_bentoml_run_command(bento_tag, '[]')
+        run_result = execute_bentoml_run_command(
+            bento_tag=bento_tag, data='[]', yatai_url=yatai_service_url
+        )
         assert 'cat' in run_result, 'Unexpected BentoService prediction result'
 
         logger.info(f'Deleting saved bundle {bento_tag}')
-        delete_svc_result = delete_bento(bento_tag)
-        assert f"{bento_tag} deleted" in delete_svc_result
+        yc.repository.delete(bento_tag=bento_tag)

e2e_tests/yatai_server/test_postgres_s3.py

@@ -1,10 +1,9 @@
 import logging
 
 from bentoml.yatai.proto.repository_pb2 import BentoUri
-from e2e_tests.cli_operations import delete_bento
+from bentoml.yatai.client import get_yatai_client
 from e2e_tests.sample_bento_service import SampleBentoService
 from e2e_tests.yatai_server.utils import (
-    get_bento_service_info,
     execute_bentoml_run_command,
     local_yatai_server,
 )
@@ -21,24 +20,26 @@ def test_yatai_server_with_postgres_and_s3(postgres_db_container_url):
 
     with local_yatai_server(
         db_url=postgres_db_container_url, repo_base_url=s3_bucket_name
-    ):
+    ) as yatai_service_url:
+        yc = get_yatai_client(yatai_service_url)
         logger.info('Saving bento service')
         svc = SampleBentoService()
-        svc.save()
+        svc.save(yatai_url=yatai_service_url)
         bento_tag = f'{svc.name}:{svc.version}'
         logger.info('BentoService saved')
 
         logger.info("Display bentoservice info")
-        bento = get_bento_service_info(svc.name, svc.version)
+        bento = yc.repository.get(bento_tag)
         logger.info(bento)
         assert (
             bento.uri.type == BentoUri.S3
         ), 'BentoService storage type mismatched, expect S3'
 
         logger.info('Validate BentoService prediction result')
-        run_result = execute_bentoml_run_command(bento_tag, '[]')
+        run_result = execute_bentoml_run_command(
+            bento_tag=bento_tag, data='[]', yatai_url=yatai_service_url
+        )
         assert 'cat' in run_result, 'Unexpected BentoService prediction result'
 
         logger.info(f'Deleting saved bundle {bento_tag}')
-        delete_svc_result = delete_bento(bento_tag)
-        assert f"{bento_tag} deleted" in delete_svc_result
+        yc.repository.delete(bento_tag=bento_tag)

e2e_tests/yatai_server/test_sqlite_gcs.py

@@ -1,13 +1,11 @@
 import logging
 
 from bentoml.yatai.proto.repository_pb2 import BentoUri
+from bentoml.yatai.client import get_yatai_client
 from e2e_tests.sample_bento_service import SampleBentoService
-from e2e_tests.cli_operations import delete_bento
 from e2e_tests.yatai_server.utils import (
     local_yatai_server,
-    get_bento_service_info,
     execute_bentoml_run_command,
-    execute_bentoml_retrieve_command,
 )
 
 logger = logging.getLogger('bentoml.test')
@@ -16,32 +14,27 @@
 def test_yatai_server_with_sqlite_and_gcs():
     gcs_bucket_name = 'gs://bentoml-e2e-tests/'
 
-    with local_yatai_server(repo_base_url=gcs_bucket_name):
+    with local_yatai_server(repo_base_url=gcs_bucket_name) as yatai_service_url:
+        yc = get_yatai_client(yatai_service_url)
         logger.info('Saving bento service')
         svc = SampleBentoService()
-        svc.save()
+        svc.save(yatai_url=yatai_service_url)
         bento_tag = f'{svc.name}:{svc.version}'
         logger.info('BentoService saved')
 
         logger.info("Display bentoservice info")
-        bento = get_bento_service_info(svc.name, svc.version)
+        bento = yc.repository.get(bento_tag)
         logger.info(bento)
         assert (
             bento.uri.type == BentoUri.GCS
         ), 'BentoService storage type mismatched, expect GCS'
 
-        retrieve_svc_result = execute_bentoml_retrieve_command(
-            f'{svc.name}:{svc.version}'
-        )
-        assert retrieve_svc_result.startswith(
-            f'Save {svc.name}:{svc.version} artifact to directory'
-        )
-
         logger.info('Validate BentoService prediction result')
-        run_result = execute_bentoml_run_command(bento_tag, '[]')
+        run_result = execute_bentoml_run_command(
+            bento_tag=bento_tag, data='[]', yatai_url=yatai_service_url
+        )
         logger.info(run_result)
         assert 'cat' in run_result, 'Unexpected BentoService prediction result'
 
         logger.info(f'Deleting saved bundle {bento_tag}')
-        delete_svc_result = delete_bento(bento_tag)
-        assert f"{bento_tag} deleted" in delete_svc_result
+        yc.repository.delete(bento_tag=bento_tag)

e2e_tests/yatai_server/test_sqlite_s3.py

@@ -1,13 +1,11 @@
 import logging
 
 from bentoml.yatai.proto.repository_pb2 import BentoUri
+from bentoml.yatai.client import get_yatai_client
 from e2e_tests.sample_bento_service import SampleBentoService
-from e2e_tests.cli_operations import delete_bento
 from e2e_tests.yatai_server.utils import (
     local_yatai_server,
-    get_bento_service_info,
     execute_bentoml_run_command,
-    execute_bentoml_retrieve_command,
 )
 
 logger = logging.getLogger('bentoml.test')
@@ -20,31 +18,27 @@ def test_yatai_server_with_sqlite_and_s3():
 
     s3_bucket_name = 's3://bentoml-e2e-test-repo/'
 
-    with local_yatai_server(repo_base_url=s3_bucket_name):
+    with local_yatai_server(repo_base_url=s3_bucket_name) as yatai_service_url:
+        yc = get_yatai_client(yatai_service_url)
         logger.info('Saving bento service')
         svc = SampleBentoService()
-        svc.save()
+        svc.save(yatai_url=yatai_service_url)
         bento_tag = f'{svc.name}:{svc.version}'
         logger.info('BentoService saved')
 
         logger.info("Display bentoservice info")
-        bento = get_bento_service_info(svc.name, svc.version)
+        bento = yc.repository.get(bento_tag)
         logger.info(bento)
         assert (
             bento.uri.type == BentoUri.S3
         ), 'BentoService storage type mismatched, expect S3'
-        retrieve_svc_result = execute_bentoml_retrieve_command(
-            f'{svc.name}:{svc.version}'
-        )
-        assert retrieve_svc_result.startswith(
-            f'Save {svc.name}:{svc.version} artifact to directory'
-        )
 
         logger.info('Validate BentoService prediction result')
-        run_result = execute_bentoml_run_command(bento_tag, '[]')
+        run_result = execute_bentoml_run_command(
+            bento_tag=bento_tag, data='[]', yatai_url=yatai_service_url
+        )
         logger.info(run_result)
         assert 'cat' in run_result, 'Unexpected BentoService prediction result'
 
         logger.info(f'Deleting saved bundle {bento_tag}')
-        delete_svc_result = delete_bento(bento_tag)
-        assert f"{bento_tag} deleted" in delete_svc_result
+        yc.repository.delete(bento_tag=bento_tag)

e2e_tests/yatai_server/utils.py

@@ -2,53 +2,27 @@
 import contextlib
 import logging
 import os
-import uuid
 
 import psutil
 
-from bentoml.utils.tempdir import TempDirectory
-from bentoml.yatai.client import YataiClient
 
 logger = logging.getLogger('bentoml.test')
 
 GRPC_PORT = '50051'
 GRPC_CHANNEL_ADDRESS = f'127.0.0.1:{GRPC_PORT}'
 
 
-def get_bento_service_info(bento_name, bento_version):
-    yatai_client = YataiClient()
-    get_result = yatai_client.repository.get(f'{bento_name}:{bento_version}')
-    return get_result
-
-
-def execute_bentoml_run_command(bento_tag, data, api="predict"):
+def execute_bentoml_run_command(bento_tag, data, api="predict", yatai_url=None):
     command = ['bentoml', 'run', bento_tag, api, '--input', data, "-q"]
+    if yatai_url is not None:
+        command.extend(['--yatai-url', yatai_url])
     proc = subprocess.Popen(
         command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=os.environ,
     )
     stdout = proc.stdout.read().decode('utf-8')
     return stdout
 
 
-def execute_bentoml_retrieve_command(bento_tag):
-    dir_name = uuid.uuid4().hex[:8]
-    with TempDirectory() as temp_dir:
-        command = [
-            'bentoml',
-            'retrieve',
-            bento_tag,
-            '--target_dir',
-            f'{temp_dir}/{dir_name}',
-        ]
-        proc = subprocess.Popen(
-            command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=os.environ
-        )
-        stdout = proc.stdout.read().decode('utf-8')
-        print(stdout)
-        print(proc.stderr.read().decode('utf-8'))
-        return stdout
-
-
 @contextlib.contextmanager
 def modified_environ(*remove, **update):
     """
@@ -100,8 +74,7 @@ def local_yatai_server(db_url=None, repo_base_url=None, port=50051):
         )
         yatai_service_url = f"localhost:{port}"
         logger.info(f'Setting config yatai_service.url to: {yatai_service_url}')
-        with modified_environ(BENTOML__YATAI_SERVICE__URL=yatai_service_url):
-            yield yatai_service_url
+        yield yatai_service_url
     finally:
         logger.info('Shutting down YataiServer gRPC server and node web server')
         kill_process(proc.pid)