Add viewer property, fix tests

wandb · Sep 16, 2021 · c85541e · c85541e
1 parent 3023cec
commit c85541e
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 1 deletion.
diff --git a/tests/test_public_api.py b/tests/test_public_api.py
@@ -522,9 +522,15 @@ def test_query_team(mock_server, api):
     assert repr(t.members[0]) == "<Member test (MEMBER)>"
 
 
+def test_viewer(mock_server, api):
+    v = api.viewer
+    assert v.admin is False
+    assert v.username == "mock"
+
+
 def test_create_service_account(mock_server, api):
     t = api.team("test")
-    assert t.create_service_account("My service account") == "Y" * 40
+    assert t.create_service_account("My service account").api_key == "Y" * 40
 
 
 def test_create_team(mock_server, api):

diff --git a/tests/utils/mock_server.py b/tests/utils/mock_server.py
@@ -543,6 +543,9 @@ def graphql():
                 "data": {
                     "viewer": {
                         "entity": "mock_server_entity",
+                        "admin": False,
+                        "email": "mock@server.test",
+                        "username": "mock",
                         "flags": '{"code_saving_enabled": true}',
                         "teams": {"edges": []},  # TODO make configurable for cli_test
                     },

diff --git a/wandb/apis/public.py b/wandb/apis/public.py
@@ -211,6 +211,8 @@ class Api(object):
             flags
             entity
             username
+            email
+            admin
             teams {
                 edges {
                     node {
@@ -268,6 +270,7 @@ def __init__(self, overrides={}, timeout: Optional[int] = None):
                 'Passing "username" to Api is deprecated. please use "entity" instead.'
             )
             self.settings["entity"] = overrides["username"]
+        self._viewer = None
         self._projects = {}
         self._runs = {}
         self._sweeps = {}
@@ -342,6 +345,15 @@ def default_entity(self):
             self._default_entity = (res.get("viewer") or {}).get("entity")
         return self._default_entity
 
+    @property
+    def viewer(self):
+        if self._viewer is None:
+            self._viewer = User(
+                self._client, self._client.execute(self.VIEWER_QUERY).get("viewer")
+            )
+            self._default_entity = self._viewer.entity
+        return self._viewer
+
     def flush(self):
         """
         The api object keeps a local cache of runs, so if the state of the run may