From 09f4133f5d717d76c5be6004483c7b4c39409f3e Mon Sep 17 00:00:00 2001 From: George Mastrapas Date: Tue, 24 May 2022 17:40:33 +0300 Subject: [PATCH 1/5] feat: load jina auth token from environment --- docarray/array/mixins/io/pushpull.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/docarray/array/mixins/io/pushpull.py b/docarray/array/mixins/io/pushpull.py index 80f626b9d0c..958469b140a 100644 --- a/docarray/array/mixins/io/pushpull.py +++ b/docarray/array/mixins/io/pushpull.py @@ -27,6 +27,14 @@ def _get_hub_config() -> Optional[Dict]: return json.load(f) +@lru_cache() +def _get_auth_token() -> Optional[str]: + _hub_config = _get_hub_config() + _config_auth_token = _hub_config.get('auth_token') if _hub_config else None + _env_auth_token = os.environ.get('JINA_AUTH_TOKEN') + return _env_auth_token or _config_auth_token + + @lru_cache() def _get_cloud_api() -> str: """Get Cloud Api for transmitting data to the cloud. @@ -73,9 +81,8 @@ def push(self, name: str, show_progress: bool = False, public: bool = True) -> D headers = {'Content-Type': ctype, **get_request_header()} - _hub_config = _get_hub_config() - if _hub_config: - auth_token = _hub_config.get('auth_token') + auth_token = _get_auth_token() + if auth_token: headers['Authorization'] = f'token {auth_token}' _head, _tail = data.split(delimiter) @@ -146,9 +153,8 @@ def pull( headers = {} - _hub_config = _get_hub_config() - if _hub_config: - auth_token = _hub_config.get('auth_token') + auth_token = _get_auth_token() + if auth_token: headers['Authorization'] = f'token {auth_token}' url = f'{_get_cloud_api()}/v2/rpc/artifact.getDownloadUrl?name={name}' From 7ea2e5d21d7f06635b77353a97bfdf1be65a1df1 Mon Sep 17 00:00:00 2001 From: George Mastrapas Date: Tue, 24 May 2022 18:52:46 +0300 Subject: [PATCH 2/5] test: add unit test for setting the token via env --- tests/unit/array/mixins/test_pushpull.py | 37 +++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/tests/unit/array/mixins/test_pushpull.py b/tests/unit/array/mixins/test_pushpull.py index 491ecf5e636..b8d773c3d82 100644 --- a/tests/unit/array/mixins/test_pushpull.py +++ b/tests/unit/array/mixins/test_pushpull.py @@ -161,9 +161,11 @@ def test_api_url_change(mocker, monkeypatch): def test_api_authorization_header(mocker, monkeypatch, tmpdir): - from docarray.array.mixins.io.pushpull import _get_hub_config + from docarray.array.mixins.io.pushpull import _get_hub_config, _get_auth_token _get_hub_config.cache_clear() + _get_auth_token.cache_clear() + os.environ['JINA_HUB_ROOT'] = str(tmpdir) token = 'test-auth-token' @@ -179,7 +181,40 @@ def test_api_authorization_header(mocker, monkeypatch, tmpdir): DocumentArray.pull(name='test_name') del os.environ['JINA_HUB_ROOT'] + + _get_hub_config.cache_clear() + _get_auth_token.cache_clear() + + assert mock.call_count == 3 # 1 for push, 1 for pull, 1 for download + + _, push_kwargs = mock.call_args_list[0] + _, pull_kwargs = mock.call_args_list[1] + + assert push_kwargs['headers'].get('Authorization') == f'token {token}' + assert pull_kwargs['headers'].get('Authorization') == f'token {token}' + + +def test_api_authorization_header_from_env(mocker, monkeypatch): + from docarray.array.mixins.io.pushpull import _get_hub_config, _get_auth_token + + _get_hub_config.cache_clear() + _get_auth_token.cache_clear() + + token = 'test-auth-token' + os.environ['JINA_AUTH_TOKEN'] = token + + mock = mocker.Mock() + _mock_post(mock, monkeypatch) + _mock_get(mock, monkeypatch) + + docs = random_docs(2) + docs.push(name='test_name') + DocumentArray.pull(name='test_name') + + del os.environ['JINA_AUTH_TOKEN'] + _get_hub_config.cache_clear() + _get_auth_token.cache_clear() assert mock.call_count == 3 # 1 for push, 1 for pull, 1 for download From 8b10cac06cd59b3ec0be68dfab1761920b9d612c Mon Sep 17 00:00:00 2001 From: George Mastrapas Date: Wed, 25 May 2022 15:51:17 +0300 Subject: [PATCH 3/5] test: add set_env_vars fixture --- tests/conftest.py | 10 ++++++++++ tests/unit/array/mixins/test_pushpull.py | 14 ++++++-------- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 09989e06156..7399317f1b8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ import tempfile import os import time +from typing import Dict import pytest @@ -33,3 +34,12 @@ def start_storage(): f"docker-compose -f {compose_yml} --project-directory . down " f"--remove-orphans" ) + + +@pytest.fixture(scope='session') +def set_env_vars(request): + _old_environ = dict(os.environ) + os.environ.update(request.param) + yield + os.environ.clear() + os.environ.update(_old_environ) diff --git a/tests/unit/array/mixins/test_pushpull.py b/tests/unit/array/mixins/test_pushpull.py index b8d773c3d82..4e88402b743 100644 --- a/tests/unit/array/mixins/test_pushpull.py +++ b/tests/unit/array/mixins/test_pushpull.py @@ -194,15 +194,15 @@ def test_api_authorization_header(mocker, monkeypatch, tmpdir): assert pull_kwargs['headers'].get('Authorization') == f'token {token}' -def test_api_authorization_header_from_env(mocker, monkeypatch): +@pytest.mark.parametrize( + 'set_env_vars', [{'JINA_AUTH_TOKEN': 'test-auth-token'}], indirect=True +) +def test_api_authorization_header_from_env(mocker, monkeypatch, set_env_vars): from docarray.array.mixins.io.pushpull import _get_hub_config, _get_auth_token _get_hub_config.cache_clear() _get_auth_token.cache_clear() - token = 'test-auth-token' - os.environ['JINA_AUTH_TOKEN'] = token - mock = mocker.Mock() _mock_post(mock, monkeypatch) _mock_get(mock, monkeypatch) @@ -211,8 +211,6 @@ def test_api_authorization_header_from_env(mocker, monkeypatch): docs.push(name='test_name') DocumentArray.pull(name='test_name') - del os.environ['JINA_AUTH_TOKEN'] - _get_hub_config.cache_clear() _get_auth_token.cache_clear() @@ -221,5 +219,5 @@ def test_api_authorization_header_from_env(mocker, monkeypatch): _, push_kwargs = mock.call_args_list[0] _, pull_kwargs = mock.call_args_list[1] - assert push_kwargs['headers'].get('Authorization') == f'token {token}' - assert pull_kwargs['headers'].get('Authorization') == f'token {token}' + assert push_kwargs['headers'].get('Authorization') == f'token test-auth-token' + assert pull_kwargs['headers'].get('Authorization') == f'token test-auth-token' From 6ee07745ff66971bb8f45d710e38db663b53a414 Mon Sep 17 00:00:00 2001 From: AlaeddineAbdessalem Date: Fri, 27 May 2022 15:37:45 +0200 Subject: [PATCH 4/5] chore: apply suggestions from code review --- docarray/array/mixins/io/pushpull.py | 8 ++++---- tests/unit/array/mixins/test_pushpull.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docarray/array/mixins/io/pushpull.py b/docarray/array/mixins/io/pushpull.py index 958469b140a..2926a30a59f 100644 --- a/docarray/array/mixins/io/pushpull.py +++ b/docarray/array/mixins/io/pushpull.py @@ -29,10 +29,10 @@ def _get_hub_config() -> Optional[Dict]: @lru_cache() def _get_auth_token() -> Optional[str]: - _hub_config = _get_hub_config() - _config_auth_token = _hub_config.get('auth_token') if _hub_config else None - _env_auth_token = os.environ.get('JINA_AUTH_TOKEN') - return _env_auth_token or _config_auth_token + hub_config = _get_hub_config() + config_auth_token = hub_config.get('auth_token') if hub_config else None + env_auth_token = os.environ.get('JINA_AUTH_TOKEN') + return env_auth_token or config_auth_token @lru_cache() diff --git a/tests/unit/array/mixins/test_pushpull.py b/tests/unit/array/mixins/test_pushpull.py index 4e88402b743..c635f5d1704 100644 --- a/tests/unit/array/mixins/test_pushpull.py +++ b/tests/unit/array/mixins/test_pushpull.py @@ -219,5 +219,5 @@ def test_api_authorization_header_from_env(mocker, monkeypatch, set_env_vars): _, push_kwargs = mock.call_args_list[0] _, pull_kwargs = mock.call_args_list[1] - assert push_kwargs['headers'].get('Authorization') == f'token test-auth-token' - assert pull_kwargs['headers'].get('Authorization') == f'token test-auth-token' + assert push_kwargs['headers'].get('Authorization') == 'token test-auth-token' + assert pull_kwargs['headers'].get('Authorization') == 'token test-auth-token' From 37542d58bc7bbc6fb49c70bfc59a5f582983fa1c Mon Sep 17 00:00:00 2001 From: Alaeddine Abdessalem Date: Fri, 27 May 2022 15:50:59 +0200 Subject: [PATCH 5/5] test: cover priority between env and config --- tests/unit/array/mixins/test_pushpull.py | 41 +++++++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/tests/unit/array/mixins/test_pushpull.py b/tests/unit/array/mixins/test_pushpull.py index c635f5d1704..c651e847aae 100644 --- a/tests/unit/array/mixins/test_pushpull.py +++ b/tests/unit/array/mixins/test_pushpull.py @@ -160,7 +160,7 @@ def test_api_url_change(mocker, monkeypatch): assert pull_kwargs['url'].startswith(test_api_url) -def test_api_authorization_header(mocker, monkeypatch, tmpdir): +def test_api_authorization_header_from_config(mocker, monkeypatch, tmpdir): from docarray.array.mixins.io.pushpull import _get_hub_config, _get_auth_token _get_hub_config.cache_clear() @@ -221,3 +221,42 @@ def test_api_authorization_header_from_env(mocker, monkeypatch, set_env_vars): assert push_kwargs['headers'].get('Authorization') == 'token test-auth-token' assert pull_kwargs['headers'].get('Authorization') == 'token test-auth-token' + + +@pytest.mark.parametrize( + 'set_env_vars', [{'JINA_AUTH_TOKEN': 'test-auth-token-env'}], indirect=True +) +def test_api_authorization_header_env_and_config( + mocker, monkeypatch, tmpdir, set_env_vars +): + from docarray.array.mixins.io.pushpull import _get_hub_config, _get_auth_token + + _get_hub_config.cache_clear() + _get_auth_token.cache_clear() + + os.environ['JINA_HUB_ROOT'] = str(tmpdir) + + token = 'test-auth-token-config' + with open(tmpdir / JINA_CLOUD_CONFIG, 'w') as f: + json.dump({'auth_token': token}, f) + + mock = mocker.Mock() + _mock_post(mock, monkeypatch) + _mock_get(mock, monkeypatch) + + docs = random_docs(2) + docs.push(name='test_name') + DocumentArray.pull(name='test_name') + + del os.environ['JINA_HUB_ROOT'] + + _get_hub_config.cache_clear() + _get_auth_token.cache_clear() + + assert mock.call_count == 3 # 1 for push, 1 for pull, 1 for download + + _, push_kwargs = mock.call_args_list[0] + _, pull_kwargs = mock.call_args_list[1] + + assert push_kwargs['headers'].get('Authorization') == 'token test-auth-token-env' + assert pull_kwargs['headers'].get('Authorization') == 'token test-auth-token-env'