fix(hubio): when error occurs, exit code should be set to non-zero (#…

…2871)

* fix(hubio): when error occurs, exit code should be set to non-zero

* fix: add mock to support hubio test

* test: fix crash test because of introducing secret.key local cache

Co-authored-by: numb3r3 <wangfelix87@gmail.com>

jina/hubble/hubio.py

@@ -156,6 +156,7 @@ def push(self) -> None:
                     f'Error while pushing session_id={req_header["jinameta-session-id"]}: '
                     f'\n{e!r}'
                 )
+                exit(1)
 
     def _prettyprint_result(self, console, result):
         # TODO: only support single executor now
@@ -369,6 +370,7 @@ def pull(self) -> str:
                     raise ValueError(f'{self.args.uri} is not a valid scheme')
             except Exception as e:
                 self.logger.error(f'{e!r}')
+                exit(2)
             finally:
                 # delete downloaded zip package if existed
                 if cached_zip_filepath is not None:

tests/unit/hubble/test_hubio.py

@@ -1,9 +1,13 @@
 import os
 import json
+from typing import NamedTuple
 import pytest
 import requests
+import itertools
+from pathlib import Path
 
-from jina.hubble.hubio import HubIO
+from jina.hubble.hubio import HubIO, HubExecutor
+from jina.hubble import helper
 from jina.parsers.hubble import set_hub_push_parser
 from jina.parsers.hubble import set_hub_pull_parser
 
@@ -37,9 +41,21 @@ def text(self):
     def status_code(self):
         return self.response_code
 
+    def iter_lines(self):
+        logs = [
+            '{"stream":"Receiving zip file..."}',
+            '{"stream":"Normalizing the content..."}',
+            '{"stream":"Building the image..."}',
+            '{"stream":"Uploading the image..."}',
+            '{"stream":"Uploading the zip..."}',
+            '{"result":{"statusCode":201,"message":"Successfully pushed w7qckiqy","data":{"executors":[{"tag":"v0","id":"w7qckiqy","image":"jinahub/w7qckiqy:v0","pullPath":"jinahub/w7qckiqy:v0","secret":"a26531e561dcb7af2c999a64cadc86d0","visibility":"public"}]}}}',
+        ]
+
+        return itertools.chain(logs)
+
 
 class GetMockResponse:
-    def __init__(self, response_code: int = 201):
+    def __init__(self, response_code: int = 200):
         self.response_code = response_code
 
     def json(self):
@@ -71,11 +87,14 @@ def status_code(self):
 def test_push(mocker, monkeypatch, path, mode):
     mock = mocker.Mock()
 
-    def _mock_post(url, files, data, headers=None):
-        mock(url=url, files=files, data=data)
+    def _mock_post(url, data, headers=None, stream=True):
+        mock(url=url, data=data)
         return PostMockResponse(response_code=requests.codes.created)
 
     monkeypatch.setattr(requests, 'post', _mock_post)
+    # Second push will use --force --secret because of .jina/secret.key
+    # Then it will use put method
+    monkeypatch.setattr(requests, 'put', _mock_post)
 
     exec_path = os.path.join(cur_dir, path)
     _args_list = [exec_path, mode]
@@ -89,7 +108,7 @@ def test_fetch(mocker, monkeypatch):
 
     def _mock_get(url, headers=None):
         mock(url=url)
-        return GetMockResponse(response_code=requests.codes.ok)
+        return GetMockResponse(response_code=200)
 
     monkeypatch.setattr(requests, 'get', _mock_get)
     args = set_hub_pull_parser().parse_args(['jinahub://dummy_mwu_encoder'])
@@ -103,11 +122,53 @@ def _mock_get(url, headers=None):
     assert executor.md5sum == 'ecbe3fdd9cbe25dbb85abaaf6c54ec4f'
 
 
-def test_pull(mocker, monkeypatch):
-    args = set_hub_pull_parser().parse_args(['jinahub://dummy_mwu_encoder'])
+class DownloadMockResponse:
+    def __init__(self, response_code: int = 200):
+        self.response_code = response_code
+
+    def iter_content(self, buffer=32 * 1024):
+
+        zip_file = Path(__file__).parent / 'dummy_executor.zip'
+        with zip_file.open('rb') as f:
+            yield f.read(buffer)
+
+    @property
+    def status_code(self):
+        return self.response_code
 
+
+def test_pull(test_envs, mocker, monkeypatch):
+    mock = mocker.Mock()
+
+    def _mock_fetch(name, tag=None, secret=None):
+        mock(name=name)
+        return HubExecutor(
+            uuid='dummy_mwu_encoder',
+            alias='alias_dummy',
+            tag='v0',
+            image_name='jinahub/pod.dummy_mwu_encoder',
+            md5sum=None,
+            visibility=True,
+            archive_url=None,
+        )
+
+    monkeypatch.setattr(HubIO, '_fetch_meta', _mock_fetch)
+
+    def _mock_download(url, stream=True, headers=None):
+        mock(url=url)
+        return DownloadMockResponse(response_code=200)
+
+    def _mock_head(url):
+        from collections import namedtuple
+
+        HeadInfo = namedtuple('HeadInfo', ['headers'])
+        return HeadInfo(headers={})
+
+    monkeypatch.setattr(requests, 'get', _mock_download)
+    monkeypatch.setattr(requests, 'head', _mock_head)
+
+    args = set_hub_pull_parser().parse_args(['jinahub://dummy_mwu_encoder'])
     HubIO(args).pull()
 
     args = set_hub_pull_parser().parse_args(['jinahub://dummy_mwu_encoder:secret'])
-
     HubIO(args).pull()