feat(hubble): add public parameter to da.push (#318)

delgermurun · web-flow · commit c9a9b35f4c4e · 2022-05-04T17:44:54.000+02:00
diff --git a/docarray/array/mixins/io/pushpull.py b/docarray/array/mixins/io/pushpull.py
@@ -57,7 +57,7 @@ class PushPullMixin:
 
     _max_bytes = 4 * 1024 * 1024 * 1024
 
-    def push(self, name: str, show_progress: bool = False) -> Dict:
+    def push(self, name: str, show_progress: bool = False, public: bool = True) -> Dict:
         """Push this DocumentArray object to Jina Cloud which can be later retrieved via :meth:`.push`
 
         .. note::
@@ -69,6 +69,7 @@ def push(self, name: str, show_progress: bool = False) -> Dict:
 
         :param name: a name that later can be used for retrieve this :class:`DocumentArray`.
         :param show_progress: if to show a progress bar on pulling
+        :param public: If True, the DocumentArray will be shared publicly. Otherwise, it will be private.
         """
         import requests
 
@@ -82,6 +83,7 @@ def push(self, name: str, show_progress: bool = False) -> Dict:
                 ),
                 'name': name,
                 'type': 'documentArray',
+                'public': public,
             }
         )
 
diff --git a/tests/unit/array/mixins/test_pushpull.py b/tests/unit/array/mixins/test_pushpull.py
@@ -1,7 +1,9 @@
+import cgi
 import json
 import os
 import pytest
 import requests
+from io import BytesIO
 
 from docarray import DocumentArray
 from docarray.array.mixins.io.pushpull import JINA_CLOUD_CONFIG
@@ -82,6 +84,27 @@ def test_push(mocker, monkeypatch):
     assert mock.call_count == 1
 
 
+@pytest.mark.parametrize('public', [True, False])
+def test_push_with_public(mocker, monkeypatch, public):
+    mock = mocker.Mock()
+    _mock_post(mock, monkeypatch)
+
+    docs = random_docs(2)
+    docs.push(name='test_name', public=public)
+
+    _, mock_kwargs = mock.call_args_list[0]
+
+    c_type, c_data = cgi.parse_header(mock_kwargs['headers']['Content-Type'])
+    assert c_type == 'multipart/form-data'
+
+    form_data = cgi.parse_multipart(
+        BytesIO(b''.join(mock_kwargs['data'])),
+        {'boundary': c_data['boundary'].encode()},
+    )
+
+    assert form_data['public'] == [str(public)]
+
+
 def test_pull(mocker, monkeypatch):
     mock = mocker.Mock()
     _mock_get(mock, monkeypatch)
@@ -104,7 +127,7 @@ def test_push_fail(mocker, monkeypatch):
     _mock_post(mock, monkeypatch, status_code=requests.codes.forbidden)
 
     docs = random_docs(2)
-    with pytest.raises(Exception) as exc_info:
+    with pytest.raises(Exception):
         docs.push('test_name')
 
     assert mock.call_count == 1
