Passing the success_exit_codes to facts (#755)

* Passing the success_exit_codes to facts.

* Fixed existing tests.

* Added test for sucess_codes.

* Minor fix from code review.

Author: mohsinhijazee

pyinfra/api/facts.py

@@ -176,6 +176,7 @@ def get_facts(
     shell_executable = state.config.SHELL
     use_sudo_password = state.config.USE_SUDO_PASSWORD
     env = state.config.ENV
+    success_exit_codes = [0]
 
     # Facts can override the shell (winrm powershell vs cmd support)
     if fact.shell_executable:
@@ -196,6 +197,7 @@ def get_facts(
         ignore_errors = current_global_kwargs.get('ignore_errors', ignore_errors)
         timeout = current_global_kwargs.get('timeout', timeout)
         env = current_global_kwargs.get('env', env)
+        success_exit_codes = current_global_kwargs.get('success_exit_codes', success_exit_codes)
 
     # Make a hash which keeps facts unique - but usable cross-deploy/threads.
     # Locks are used to maintain order.
@@ -243,6 +245,7 @@ def get_facts(
                 sudo=sudo,
                 sudo_user=sudo_user,
                 use_sudo_password=use_sudo_password,
+                success_exit_codes=success_exit_codes,
                 su_user=su_user,
                 timeout=timeout,
                 env=env,

tests/test_api/test_api_facts.py

@@ -33,6 +33,7 @@ def test_get_fact(self):
             print_input=False,
             print_output=False,
             shell_executable='sh',
+            success_exit_codes=[0],
             su_user=None,
             sudo=False,
             sudo_user=None,
@@ -72,6 +73,7 @@ def test_get_fact_current_op_meta(self):
             print_input=False,
             print_output=False,
             shell_executable='sh',
+            success_exit_codes=[0],
             su_user='someuser',
             sudo=True,
             sudo_user='someuser',
@@ -104,6 +106,7 @@ def test_get_fact_error(self):
             print_input=False,
             print_output=False,
             shell_executable='sh',
+            success_exit_codes=[0],
             su_user=None,
             sudo=False,
             sudo_user=None,
@@ -143,6 +146,39 @@ def test_get_fact_error_ignore(self):
             print_input=False,
             print_output=False,
             shell_executable='sh',
+            success_exit_codes=[0],
+            su_user=None,
+            sudo=False,
+            sudo_user=None,
+            timeout=None,
+            env={},
+            use_sudo_password=False,
+            return_combined_output=True,
+        )
+
+    def test_get_fact_with_list_of_success_codes(self):
+        inventory = make_inventory(hosts=('anotherhost',))
+        state = State(inventory, Config())
+
+        anotherhost = inventory.get_host('anotherhost')
+
+        connect_all(state)
+
+        exit_codes = [1, 2, 3]
+        with patch('pyinfra.api.connectors.ssh.run_shell_command') as fake_run_command:
+            fake_run_command.return_value = MagicMock(), [('stdout', 'some-output')]
+            fact_data = get_facts(state, 'command', ('yes',), {'success_exit_codes': exit_codes})
+
+        assert fact_data == {anotherhost: 'some-output'}
+
+        fake_run_command.assert_called_with(
+            state,
+            anotherhost,
+            'yes',
+            print_input=False,
+            print_output=False,
+            shell_executable='sh',
+            success_exit_codes=exit_codes,
             su_user=None,
             sudo=False,
             sudo_user=None,