Add protected parameter to branches method with simple unit test.

bboe · bboe · commit 96b48c7fff57 · 2015-11-15T09:13:16.000-08:00
diff --git a/github3/repos/repo.py b/github3/repos/repo.py
@@ -399,18 +399,27 @@ def branch(self, name):
             json = self._json(self._get(url), 200)
         return self._instance_or_null(Branch, json)
 
-    def branches(self, number=-1, etag=None):
+    def branches(self, number=-1, protected=False, etag=None):
         r"""Iterate over the branches in this repository.
 
         :param int number: (optional), number of branches to return. Default:
             -1 returns all branches
+        :param bool protected: (optional), True lists only protected branches.
+            Default: False
         :param str etag: (optional), ETag from a previous request to the same
             endpoint
         :returns: generator of
             :class:`Branch <github3.repos.branch.Branch>`\ es
         """
         url = self._build_url('branches', base_url=self._api)
-        return self._iter(int(number), url, Branch, etag=etag)
+
+        # The Accept header will likely be removable once the feature is out of
+        # preview mode. See: http://git.io/v4O1e
+        headers = {'Accept': 'application/vnd.github.loki-preview+json'}
+
+        params = {'protected': '1'} if protected else None
+        return self._iter(int(number), url, Branch, params, etag=etag,
+                          headers=headers)
 
     def code_frequency(self, number=-1, etag=None):
         """Iterate over the code frequency per week.
diff --git a/tests/unit/test_repos_repo.py b/tests/unit/test_repos_repo.py
@@ -368,7 +368,18 @@ def test_branches(self):
         self.session.get.assert_called_once_with(
             url_for('branches'),
             params={'per_page': 100},
-            headers={}
+            headers={'Accept': 'application/vnd.github.loki-preview+json'}
+        )
+
+    def test_branches_protected(self):
+        """Test ability to iterate over protected branches in a Repository."""
+        i = self.instance.branches(protected=True)
+        self.get_next(i)
+
+        self.session.get.assert_called_once_with(
+            url_for('branches'),
+            params={'per_page': 100, 'protected': '1'},
+            headers={'Accept': 'application/vnd.github.loki-preview+json'}
         )
 
     def test_code_frequency(self):
