From c78ebd0d85fcf5edb5fa71966733b10ca0ef62bd Mon Sep 17 00:00:00 2001
From: Evan J Felix <e@pnnl.gov>
Date: Fri, 17 Mar 2017 13:10:51 -0700
Subject: [PATCH] add Tests and cleanup testing issues

---
 .pre-commit-hooks.yaml                  |  6 +++++
 README.md                               |  2 ++
 pre_commit_hooks/no_commit_to_branch.py | 17 +++++++++-----
 tests/check_no_commit_to_branch_test.py | 30 +++++++++++++++++++++++++
 4 files changed, 49 insertions(+), 6 deletions(-)
 create mode 100644 tests/check_no_commit_to_branch_test.py

diff --git a/.pre-commit-hooks.yaml b/.pre-commit-hooks.yaml
index 0604585..f7e35d4 100644
--- a/.pre-commit-hooks.yaml
+++ b/.pre-commit-hooks.yaml
@@ -147,3 +147,9 @@
     entry: trailing-whitespace-fixer
     language: python
     files: \.(asciidoc|adoc|coffee|cpp|css|c|ejs|erb|groovy|h|haml|hh|hpp|hxx|html|in|j2|jade|json|js|less|markdown|md|ml|mli|pp|py|rb|rs|R|scala|scss|sh|slim|tex|tmpl|ts|txt|yaml|yml)$
+-   id: no-commit-to-branch
+    name: Dont commit to branch
+    entry: no-commit-to-branch
+    language: python
+    files: .*
+    always_run: true
diff --git a/README.md b/README.md
index 031edfd..723136d 100644
--- a/README.md
+++ b/README.md
@@ -56,6 +56,8 @@ Add this to your `.pre-commit-config.yaml`
 - `forbid-new-submodules` - Prevent addition of new git submodules.
 - `name-tests-test` - Assert that files in tests/ end in `_test.py`.
     - Use `args: ['--django']` to match `test*.py` instead.
+- `no-commit-to-branch` - Protect specific branches from direct checkins.
+    - Use `args: -b <branch> ` to set the branch. `master` is the default if no argument is set.
 - `pyflakes` - Run pyflakes on your python files.
 - `pretty-format-json` - Checks that all your JSON files are pretty.  "Pretty"
   here means that keys are sorted and indented.  You can configure this with
diff --git a/pre_commit_hooks/no_commit_to_branch.py b/pre_commit_hooks/no_commit_to_branch.py
index dd88078..d4170ed 100644
--- a/pre_commit_hooks/no_commit_to_branch.py
+++ b/pre_commit_hooks/no_commit_to_branch.py
@@ -5,6 +5,13 @@ import sys
 
 import util
 
+def is_on_branch(protected):
+    retval = False
+    branch = util.cmd_output('git', 'symbolic-ref', 'HEAD')
+    chunks = branch.strip().split('/')
+    if chunks[2] == protected:
+        retval = True
+    return retval
 
 def main(argv=None):
     parser = argparse.ArgumentParser()
@@ -12,12 +19,10 @@ def main(argv=None):
     parser.add_argument('filenames', nargs='*', help='filenames to check.')
     args = parser.parse_args(argv)
 
-    retval = 0
-    branch = util.cmd_output('git', 'symbolic-ref', 'HEAD')
-    chunks = branch.strip().split('/')
-    if chunks[2] == args.b:
-        retval = -1
-    return retval
+    if is_on_branch(args.b):
+        return 1
+    else:
+        return 0
 
 
 if __name__ == '__main__':
diff --git a/tests/check_no_commit_to_branch_test.py b/tests/check_no_commit_to_branch_test.py
new file mode 100644
index 0000000..e95d445
--- /dev/null
+++ b/tests/check_no_commit_to_branch_test.py
@@ -0,0 +1,30 @@
+from __future__ import absolute_import
+from __future__ import unicode_literals
+
+import subprocess
+import sys
+import pytest
+
+from pre_commit_hooks.no_commit_to_branch import is_on_branch
+from pre_commit_hooks.no_commit_to_branch import main
+from pre_commit_hooks.util import cmd_output
+
+
+def test_other_branch(temp_git_dir):
+    with temp_git_dir.as_cwd():
+        cmd_output('git', 'checkout', '-b', 'anotherbranch')
+        assert is_on_branch('master') == False
+
+def test_master_branch(temp_git_dir):
+    with temp_git_dir.as_cwd():
+        assert is_on_branch('master') == True
+
+def test_main_other_call(temp_git_dir):
+    with temp_git_dir.as_cwd():
+        cmd_output('git', 'checkout', '-b', 'other')
+        assert main(['-b','other']) == 1
+
+def test_main_default_call(temp_git_dir):
+    with temp_git_dir.as_cwd():
+        cmd_output('git', 'checkout', '-b', 'anotherbranch')
+        assert main() == 0