From db4b8f02f3a9e833d579747a6d0d56a2d31d36f8 Mon Sep 17 00:00:00 2001
From: Guy Kisel <gkisel@riotgames.com>
Date: Wed, 11 Mar 2015 17:44:59 -0700
Subject: [PATCH] Add --django flag to test name checker

---
 pre_commit_hooks/tests_should_end_in_test.py | 16 ++++++++++++++--
 tests/tests_should_end_in_test_test.py       | 20 ++++++++++++++++++++
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/pre_commit_hooks/tests_should_end_in_test.py b/pre_commit_hooks/tests_should_end_in_test.py
index 14c672e..5df6fd8 100644
--- a/pre_commit_hooks/tests_should_end_in_test.py
+++ b/pre_commit_hooks/tests_should_end_in_test.py
@@ -1,23 +1,35 @@
 from __future__ import print_function
 
 import argparse
+import re
 import sys
 
 
 def validate_files(argv=None):
     parser = argparse.ArgumentParser()
     parser.add_argument('filenames', nargs='*')
+    parser.add_argument(
+        '--django', default=False, action='store_true',
+        help='Use Django-style test naming pattern (test*.py)'
+    )
     args = parser.parse_args(argv)
 
     retcode = 0
+    test_name_pattern = '.*_test.py'
+    if args.django:
+        test_name_pattern = 'test.*.py'
     for filename in args.filenames:
         if (
-                not filename.endswith('_test.py') and
+                not re.match(test_name_pattern, filename) and
                 not filename.endswith('__init__.py') and
                 not filename.endswith('/conftest.py')
         ):
             retcode = 1
-            print('{0} does not end in _test.py'.format(filename))
+            print(
+                '{0} does not match pattern "{1}"'.format(
+                    filename, test_name_pattern
+                )
+            )
 
     return retcode
 
diff --git a/tests/tests_should_end_in_test_test.py b/tests/tests_should_end_in_test_test.py
index 3ba1617..ba27b19 100644
--- a/tests/tests_should_end_in_test_test.py
+++ b/tests/tests_should_end_in_test_test.py
@@ -9,3 +9,23 @@ def test_validate_files_all_pass():
 def test_validate_files_one_fails():
     ret = validate_files(['not_test_ending.py', 'foo_test.py'])
     assert ret == 1
+
+
+def test_validate_files_django_all_pass():
+    ret = validate_files(['--django', 'test_foo.py', 'test_bar.py'])
+    assert ret == 0
+
+
+def test_validate_files_django_one_fails():
+    ret = validate_files(['--django', 'not_test_ending.py', 'test_foo.py'])
+    assert ret == 1
+
+
+def test_validate_files_not_django_fails():
+    ret = validate_files(['foo_test.py', 'bar_test.py', 'test_baz.py'])
+    assert ret == 1
+
+
+def test_validate_files_django_fails():
+    ret = validate_files(['--django', 'foo_test.py', 'test_bar.py', 'test_baz.py'])
+    assert ret == 1