feat(client): adapt with adding returing data "branchName" in "listDr…

…afts"

PR Closed: #645

tensorbay/client/dataset.py

@@ -443,9 +443,9 @@ def checkout(self, revision: Optional[str] = None, draft_number: Optional[int] =
                 self._status.branch_name = None
 
         if draft_number:
-            draft_number = self.get_draft(draft_number).number
-            self._status.checkout(draft_number=draft_number)
-            self._status.branch_name = None
+            draft = self.get_draft(draft_number)
+            self._status.checkout(draft_number=draft.number)
+            self._status.branch_name = draft.branch_name
 
     def commit(self, message: str, *, tag: Optional[str] = None) -> None:
         """Commit the draft.

tensorbay/client/struct.py

@@ -279,23 +279,26 @@ class Draft(ReprMixin, EqMixin):
     Arguments:
         number: The number of the draft.
         title: The title of the draft.
+        branch_name: The branch name.
 
     """
 
     _T = TypeVar("_T", bound="Draft")
 
     _repr_attrs = ("title",)
 
-    def __init__(self, number: int, title: str) -> None:
+    def __init__(self, number: int, title: str, branch_name: str) -> None:
         self.number = number
         self.title = title
+        self.branch_name = branch_name
 
     def _repr_head(self) -> str:
         return f"{self.__class__.__name__}({self.number})"
 
     def _loads(self, contents: Dict[str, Any]) -> None:
         self.number = contents["number"]
         self.title = contents["title"]
+        self.branch_name = contents["branchName"]
 
     @classmethod
     def loads(cls: Type[_T], contents: Dict[str, Any]) -> _T:
@@ -307,6 +310,7 @@ def loads(cls: Type[_T], contents: Dict[str, Any]) -> _T:
                     {
                         "number": <int>
                         "title": <str>
+                        "branchName": <str>
                     }
 
         Returns:
@@ -324,7 +328,8 @@ def dumps(self) -> Dict[str, Any]:
                 {
                     "number": <int>
                     "title": <str>
+                    "branchName": <str>
                 }
 
         """
-        return {"number": self.number, "title": self.title}
+        return {"number": self.number, "title": self.title, "branchName": self.branch_name}

tensorbay/client/tests/test_struct.py

@@ -26,7 +26,8 @@
 
 _DRAFT_NUMBER = 1
 _DRAFT_TITLE = "draft title"
-_DRAFT_DATA = {"number": _DRAFT_NUMBER, "title": _DRAFT_TITLE}
+_DRAFT_BRANCH_NAME = "main"
+_DRAFT_DATA = {"number": _DRAFT_NUMBER, "title": _DRAFT_TITLE, "branchName": _DRAFT_BRANCH_NAME}
 
 
 class TestUser:
@@ -93,15 +94,16 @@ def test_dumps(self):
 
 class TestDraft:
     def test_init(self):
-        draft = Draft(_DRAFT_NUMBER, _DRAFT_TITLE)
+        draft = Draft(_DRAFT_NUMBER, _DRAFT_TITLE, _DRAFT_BRANCH_NAME)
         assert draft.number == _DRAFT_NUMBER
         assert draft.title == _DRAFT_TITLE
 
     def test_loads(self):
         draft = Draft.loads(_DRAFT_DATA)
         assert draft.number == _DRAFT_DATA["number"]
         assert draft.title == _DRAFT_DATA["title"]
+        assert draft.branch_name == _DRAFT_DATA["branchName"]
 
     def test_dumps(self):
-        draft = Draft(_DRAFT_NUMBER, _DRAFT_TITLE)
+        draft = Draft(_DRAFT_NUMBER, _DRAFT_TITLE, _DRAFT_BRANCH_NAME)
         assert draft.dumps() == _DRAFT_DATA

tests/test_dataset.py

@@ -190,6 +190,7 @@ def test_checkout(self, accesskey, url):
         with pytest.raises(ResourceNotExistError):
             dataset_client.checkout(draft_number=2)
         dataset_client.checkout(draft_number=3)
+        assert dataset_client._status.branch_name == DEFAULT_BRANCH
         assert dataset_client._status.draft_number == 3
 
         gas_client.delete_dataset(dataset_name)

tests/test_draft.py

@@ -6,6 +6,7 @@
 import pytest
 
 from tensorbay.client import GAS
+from tensorbay.client.gas import DEFAULT_BRANCH
 from tensorbay.client.struct import Draft
 from tensorbay.exception import CommitStatusError, ResponseError
 
@@ -44,7 +45,7 @@ def test_list_drafts(self, accesskey, url):
 
         drafts = dataset_client.list_drafts()
         assert len(drafts) == 1
-        assert drafts[0] == Draft(draft_number_2, "draft-2")
+        assert drafts[0] == Draft(draft_number_2, "draft-2", DEFAULT_BRANCH)
 
         with pytest.raises(TypeError):
             get_draft_number_by_title(dataset_client.list_drafts(), "draft-3")