diff --git a/tests/rules/test_apt_get.py b/tests/rules/test_apt_get.py index 56ad8208..ca2d95d9 100644 --- a/tests/rules/test_apt_get.py +++ b/tests/rules/test_apt_get.py @@ -16,6 +16,8 @@ def test_match(command): @pytest.mark.parametrize('command, return_value', [ (Command(script='vim', stderr='vim: command not found'), + [('vim', 'main'), ('vim-tiny', 'main')]), + (Command(script='sudo vim', stderr='vim: command not found'), [('vim', 'main'), ('vim-tiny', 'main')])]) @patch('thefuck.rules.apt_get.CommandNotFound', create=True) @patch.multiple(apt_get, create=True, apt_get='apt_get') @@ -38,7 +40,9 @@ def test_not_match(command): reason='Skip if python-commandnotfound is not available') @pytest.mark.parametrize('command, new_command', [ (Command('vim'), 'sudo apt-get install vim && vim'), - (Command('convert'), 'sudo apt-get install imagemagick && convert')]) + (Command('convert'), 'sudo apt-get install imagemagick && convert'), + (Command('sudo vim'), 'sudo apt-get install vim && sudo vim'), + (Command('sudo convert'), 'sudo apt-get install imagemagick && sudo convert')]) def test_get_new_command(command, new_command): assert get_new_command(command, None) == new_command @@ -47,6 +51,11 @@ def test_get_new_command(command, new_command): (Command('vim'), 'sudo apt-get install vim && vim', [('vim', 'main'), ('vim-tiny', 'main')]), (Command('convert'), 'sudo apt-get install imagemagick && convert', + [('imagemagick', 'main'), + ('graphicsmagick-imagemagick-compat', 'universe')]), + (Command('sudo vim'), 'sudo apt-get install vim && sudo vim', + [('vim', 'main'), ('vim-tiny', 'main')]), + (Command('sudo convert'), 'sudo apt-get install imagemagick && sudo convert', [('imagemagick', 'main'), ('graphicsmagick-imagemagick-compat', 'universe')])]) @patch('thefuck.rules.apt_get.CommandNotFound', create=True) @@ -55,5 +64,3 @@ def test_get_new_command_mocked(cmdnf_mock, command, new_command, return_value): get_packages = Mock(return_value=return_value) cmdnf_mock.CommandNotFound.return_value = Mock(getPackages=get_packages) assert get_new_command(command, None) == new_command - assert cmdnf_mock.CommandNotFound.called - assert get_packages.called diff --git a/thefuck/rules/apt_get.py b/thefuck/rules/apt_get.py index 1b727188..7dc306f2 100644 --- a/thefuck/rules/apt_get.py +++ b/thefuck/rules/apt_get.py @@ -1,27 +1,30 @@ from thefuck import shells -from thefuck.utils import sudo_support +from thefuck.utils import memoize try: import CommandNotFound except ImportError: enabled_by_default = False -@sudo_support -def match(command, settings): - if 'not found' in command.stderr: - try: - c = CommandNotFound.CommandNotFound() - pkgs = c.getPackages(command.script.split(" ")[0]) - name, _ = pkgs[0] - return True - except IndexError: - # IndexError is thrown when no matching package is found - return False -@sudo_support +@memoize +def get_package(command): + try: + c = CommandNotFound.CommandNotFound() + cmd = command.split(' ') + pkgs = c.getPackages(cmd[0] if cmd[0] != 'sudo' else cmd[1]) + name, _ = pkgs[0] + return name + except IndexError: + # IndexError is thrown when no matching package is found + return None + + +def match(command, settings): + return 'not found' in command.stderr and get_package(command.script) + + def get_new_command(command, settings): - c = CommandNotFound.CommandNotFound() - pkgs = c.getPackages(command.script.split(" ")[0]) - name, _ = pkgs[0] + name = get_package(command.script) formatme = shells.and_('sudo apt-get install {}', '{}') return formatme.format(name, command.script)