client: better init and pbar (#977)

mpcontribs-client/mpcontribs/client/__init__.py

@@ -76,9 +76,6 @@
 SUPPORTED_FILETYPES = (Gz, Jpeg, Png, Gif, Tiff)
 SUPPORTED_MIMES = [t().mime for t in SUPPORTED_FILETYPES]
 DEFAULT_DOWNLOAD_DIR = Path.home() / "mpcontribs-downloads"
-EMPTY_SPEC_DICT = {
-    "swagger": "2.0", "paths": {}, "info": {"title": "Swagger", "version": "0.0"},
-}
 
 j2h = Json2Html()
 pd.options.plotting.backend = "plotly"
@@ -238,7 +235,7 @@ def _response_hook(resp, *args, **kwargs):
         resp.result = resp.content
         resp.count = 1
     else:
-        logger.error(resp.status_code)
+        logger.error(f"request failed with status {resp.status_code}!")
         resp.count = 0
 
 
@@ -393,14 +390,16 @@ def from_data(cls, name: str, data: Union[list, dict]):
         )
 
 
-def _run_futures(futures, total: int = 0, timeout: int = -1, desc=None):
+def _run_futures(futures, total: int = 0, timeout: int = -1, desc=None, disable=False):
     """helper to run futures/requests"""
     start = time.perf_counter()
     total_set = total > 0
     total = total if total_set else len(futures)
     responses = {}
 
-    with tqdm(total=total, desc=desc, file=tqdm_out) as pbar:
+    with tqdm(
+        total=total, desc=desc, file=tqdm_out, miniters=1, delay=5, disable=disable
+    ) as pbar:
         for future in as_completed(futures):
             if not future.cancelled():
                 response = future.result()
@@ -433,26 +432,38 @@ def _load(protocol, host, headers_json, project):
     origin_url = f"{url}/apispec.json"
     fn = urlsafe_b64encode(origin_url.encode('utf-8')).decode('utf-8')
     apispec = Path(gettempdir()) / fn
+    spec_dict = None
 
     if apispec.exists():
         spec_dict = ujson.loads(apispec.read_bytes())
         logger.debug(f"Specs for {origin_url} re-loaded from {apispec}.")
     else:
-        try:
-            if requests.options(f"{url}/healthcheck").status_code == 200:
-                loader = Loader(http_client)
-                spec_dict = loader.load_spec(origin_url)
+        retries, max_retries = 0, 3
+        while retries < max_retries:
+            try:
+                is_mock_test = 'unittest' in sys.modules and protocol == "http"
+                if is_mock_test or requests.options(f"{url}/healthcheck").status_code == 200:
+                    loader = Loader(http_client)
+                    spec_dict = loader.load_spec(origin_url)
 
-                with apispec.open("w") as f:
-                    ujson.dump(spec_dict, f)
+                    with apispec.open("w") as f:
+                        ujson.dump(spec_dict, f)
 
-                logger.debug(f"Specs for {origin_url} saved as {apispec}.")
-            else:
-                spec_dict = EMPTY_SPEC_DICT
-                logger.error(f"Specs not loaded: Healthcheck for {url} failed!")
-        except RequestException:
-            spec_dict = EMPTY_SPEC_DICT
-            logger.error(f"Specs not loaded: Could not connect to {url}!")
+                    logger.debug(f"Specs for {origin_url} saved as {apispec}.")
+                    break
+                else:
+                    retries += 1
+                    logger.warning(
+                        f"Specs not loaded: Healthcheck for {url} failed! Waiting 60s ..."
+                    )
+                    time.sleep(60)
+            except RequestException:
+                retries += 1
+                logger.warning(f"Specs not loaded: Could not connect to {url}! Waiting 60s ...")
+                time.sleep(60)
+
+    if not spec_dict:
+        raise ValueError(f"Could not load specs from {url}!")  # not cached
 
     spec_dict["host"] = host
     spec_dict["schemes"] = [protocol]
@@ -480,11 +491,11 @@ def _load(protocol, host, headers_json, project):
     future = session.get(f"{url}/projects/", **kwargs)
     track_id = "get_columns"
     setattr(future, "track_id", track_id)
-    resp = _run_futures([future], timeout=3).get(track_id, {}).get("result")
+    resp = _run_futures([future], timeout=3, disable=True).get(track_id, {}).get("result")
     session.close()
 
     if not resp or not resp["data"]:
-        return swagger_spec
+        raise ValueError(f"Failed to load projects for query {query}!")
 
     if project and not resp["data"]:
         raise ValueError(f"{project} doesn't exist, or access denied!")

mpcontribs-client/mpcontribs/client/test_client.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 import pytest
 
+from unittest.mock import patch, MagicMock
 from mpcontribs.client import validate_email, Client, DEFAULT_HOST, email_format
 from swagger_spec_validator.common import SwaggerValidationError
 
@@ -12,48 +13,59 @@ def test_validate_email():
         validate_email("fake:[email protected]")
 
 
-def test_Client():
-    kwargs = {"apikey": "1234"}
-    spec = Client(**kwargs).swagger_spec
-    assert spec.http_client.headers == {
-        "Content-Type": "application/json", "x-api-key": "1234"
-    }
-    assert spec.origin_url == f"https://{DEFAULT_HOST}/apispec.json"
-    assert spec.spec_dict["host"] == DEFAULT_HOST
-    assert spec.spec_dict["schemes"] == ["https"]
-    assert spec.user_defined_formats["email"] == email_format
-
-    kwargs = {"headers": {"a": "b"}, "host": "localhost:10000"}
-    spec = Client(**kwargs).swagger_spec
-    assert spec.http_client.headers == {
-        "Content-Type": "application/json", "a": "b"
-    }
-    assert spec.origin_url == "http://localhost:10000/apispec.json"
-    assert spec.spec_dict["host"] == "localhost:10000"
-    assert spec.spec_dict["schemes"] == ["http"]
-    assert spec.user_defined_formats["email"] == email_format
-
-    kwargs = {"host": "contribs-apis:10000"}
-    spec = Client(**kwargs).swagger_spec
-    assert spec.http_client.headers == {"Content-Type": "application/json"}
-    assert spec.origin_url == "http://contribs-apis:10000/apispec.json"
-    assert spec.spec_dict["host"] == "contribs-apis:10000"
-    assert spec.spec_dict["schemes"] == ["http"]
-    assert spec.user_defined_formats["email"] == email_format
-
-    kwargs = {"host": "ml-api.materialsproject.org"}
-    spec = Client(**kwargs).swagger_spec
-    assert spec.http_client.headers == {"Content-Type": "application/json"}
-    assert spec.origin_url == "https://ml-api.materialsproject.org/apispec.json"
-    assert spec.spec_dict["host"] == "ml-api.materialsproject.org"
-    assert spec.spec_dict["schemes"] == ["https"]
-    assert spec.user_defined_formats["email"] == email_format
+@patch(
+    "bravado.swagger_model.Loader.load_spec",
+    new=MagicMock(
+        return_value={
+            "swagger": "2.0",
+            "paths": {},
+            "info": {"title": "Swagger", "version": "0.0"},
+        }
+    ),
+)
+def test_mock():
+    host = "localhost:10000"
+    with Client(host=host, headers={"a": "b"}) as client:
+        spec = client.swagger_spec
+        assert spec.http_client.headers == {
+            "Content-Type": "application/json", "a": "b"
+        }
+        assert spec.origin_url == f"http://{host}/apispec.json"
+        assert spec.spec_dict["host"] == host
+        assert spec.spec_dict["schemes"] == ["http"]
+        assert spec.user_defined_formats["email"] == email_format
+
+    host = "contribs-apis:10000"
+    with Client(host=host) as client:
+        spec = client.swagger_spec
+        assert spec.http_client.headers == {"Content-Type": "application/json"}
+        assert spec.origin_url == f"http://{host}/apispec.json"
+        assert spec.spec_dict["host"] == host
+        assert spec.spec_dict["schemes"] == ["http"]
+        assert spec.user_defined_formats["email"] == email_format
 
     with pytest.raises(ValueError):
-        kwargs = {"host": "not.valid.org"}
-        spec = Client(**kwargs).swagger_spec
+        with Client(host="not.valid.org") as client:
+            spec = client.swagger_spec
 
 
-def test_Client_Live():
-    with Client() as client:
+def test_live():
+    with Client(apikey="1234") as client:
         assert client.url == f"https://{DEFAULT_HOST}"
+        spec = client.swagger_spec
+        assert spec.http_client.headers == {
+            "Content-Type": "application/json", "x-api-key": "1234"
+        }
+        assert spec.origin_url == f"https://{DEFAULT_HOST}/apispec.json"
+        assert spec.spec_dict["host"] == DEFAULT_HOST
+        assert spec.spec_dict["schemes"] == ["https"]
+        assert spec.user_defined_formats["email"] == email_format
+
+    host = "ml-api.materialsproject.org"
+    with Client(host=host) as client:
+        spec = client.swagger_spec
+        assert spec.http_client.headers == {"Content-Type": "application/json"}
+        assert spec.origin_url == f"https://{host}/apispec.json"
+        assert spec.spec_dict["host"] == host
+        assert spec.spec_dict["schemes"] == ["https"]
+        assert spec.user_defined_formats["email"] == email_format