use fixtures and regex

tests/rules/test_git_lfs_mistype.py

@@ -1,9 +1,12 @@
+import pytest
+
 from thefuck.rules.git_lfs_mistype import match, get_new_command
 from thefuck.types import Command
 
 
-def test_match():
-    err_response = """
+@pytest.fixture
+def mistype_response():
+    return """
 Error: unknown command "evn" for "git-lfs"
 
 Did you mean this?
@@ -12,51 +15,15 @@ Did you mean this?
 
 Run 'git-lfs --help' for usage.
     """
-    assert match(Command('git lfs evn', err_response))
 
 
-def test_not_match():
+def test_match(mistype_response):
+    assert match(Command('git lfs evn', mistype_response))
     err_response = 'bash: git: command not found'
     assert not match(Command('git lfs env', err_response))
+    assert not match(Command('docker lfs env', mistype_response))
 
 
-def test_not_git_command():
-    err_response = """
-Error: unknown command "evn" for "git-lfs"
-
-Did you mean this?
-        env
-        ext
-
-Run 'git-lfs --help' for usage.
-    """
-    assert not match(Command('docker lfs env', err_response))
-
-
-def test_get_new_command():
-    err_response = """
-Error: unknown command "evn" for "git-lfs"
-
-Did you mean this?
-        env
-        ext
-
-Run 'git-lfs --help' for usage.
-    """
-    result = get_new_command(Command('git lfs evn', err_response))
-    expected = 'git lfs env'
-    assert result == expected
-
-
-def test_get_another_new_command():
-    err_response = """
-Error: unknown command "chekout" for "git-lfs"
-
-Did you mean this?
-        checkout
-
-Run 'git-lfs --help' for usage.
-    """
-    result = get_new_command(Command('git lfs chekout', err_response))
-    expected = 'git lfs checkout'
-    assert result == expected
+def test_get_new_command(mistype_response):
+    assert (get_new_command(Command('git lfs evn', mistype_response))
+            == ['git lfs env', 'git lfs ext'])

thefuck/rules/git_lfs_mistype.py

@@ -1,4 +1,5 @@
-from thefuck.shells import shell
+import re
+from thefuck.utils import get_all_matched_commands, replace_command
 from thefuck.specific.git import git_support
 
 
@@ -12,7 +13,6 @@ def match(command):
 
 @git_support
 def get_new_command(command):
-    new = command.script.split(' ')
-    recommended = command.output.split('\n')[4].strip()
-    new[2] = recommended
-    return shell.and_(' '.join(new))
+    broken_cmd = re.findall(r'Error: unknown command "([^"]*)" for "git-lfs"', command.output)[0]
+    matched = get_all_matched_commands(command.output, ['Did you mean'])
+    return replace_command(command, broken_cmd, matched)