Wrong order of operation when using a context manager and python function

Describe the bug

For some reason, having some operations run from a context manager and some other encapsulated in a python function raise a wrong order of operations. It is not clear yet if this is due to the context manager or not.

See non working and working example below, notice the only difference is the fact that the command is put in a python function.

To Reproduce

DOES NOT WORK:

import os
from pyinfra import host
from pyinfra.operations import files, server


class TmpDir:
    def __init__(self, name):
        self.tmp_dir_name = f"tmp_{name}"
        self.tmp_dir_path = os.path.join(os.sep, "home", host.data.ssh_user, "tmp", self.tmp_dir_name)
        self.create_tmp_dir = None

    def __enter__(self):
        self.create_tmp_dir = files.directory(name=f"Create temporary directory [{self.tmp_dir_name}]", path=self.tmp_dir_path, recursive=True)
        return self.tmp_dir_path

    def __exit__(self, type, value, traceback):
        files.directory(name=f"Destroy temporary directory [{self.tmp_dir_name}]", path=self.tmp_dir_path, present=False, assume_present=self.create_tmp_dir.changed)


def run_command(path):
    server.shell(name="Test ls", commands="ls .", _chdir=path)


with TmpDir("test") as tmp_dir_path:
    run_command(tmp_dir_path)

--> Loading config...
--> Loading inventory...

--> Connecting to hosts...
    No host key for *** found in known_hosts
    [***] Connected

--> Preparing operations...
    Loading: tasks/cm.py
    [***] Ready: tasks/cm.py

--> Proposed changes:
    Groups: ** / ***
    [***]   Operations: 3   Commands: 3

--> Beginning operation run...
--> Starting operation: Create temporary directory [tmp_test]
    [***] Success

--> Starting operation: Destroy temporary directory [tmp_test]
    [***] Success

--> Starting operation: Test ls
    [***] sh: line 0: cd: /home/***/tmp/tmp_test: No such file or directory
    [***] Error
No hosts remaining!
--> pyinfra error: 

WORKS:

import os
from pyinfra import host
from pyinfra.operations import files, server


class TmpDir:
    def __init__(self, name):
        self.tmp_dir_name = f"tmp_{name}"
        self.tmp_dir_path = os.path.join(os.sep, "home", host.data.ssh_user, "tmp", self.tmp_dir_name)
        self.create_tmp_dir = None

    def __enter__(self):
        self.create_tmp_dir = files.directory(name=f"Create temporary directory [{self.tmp_dir_name}]", path=self.tmp_dir_path, recursive=True)
        return self.tmp_dir_path

    def __exit__(self, type, value, traceback):
        files.directory(name=f"Destroy temporary directory [{self.tmp_dir_name}]", path=self.tmp_dir_path, present=False, assume_present=self.create_tmp_dir.changed)

with TmpDir("test") as tmp_dir_path:
    server.shell(name="Test ls", commands="ls .", _chdir=tmp_dir_path)

--> Loading config...
--> Loading inventory...

--> Connecting to hosts...
    No host key for *** found in known_hosts
    [***] Connected

--> Preparing operations...
    Loading: tasks/cm.py
    [***] Ready: tasks/cm.py

--> Proposed changes:
    Groups: ** / **
    [***]   Operations: 3   Commands: 3

--> Beginning operation run...
--> Starting operation: Create temporary directory [tmp_test]
    [***] Success

--> Starting operation: Test ls
    [***] Success

--> Starting operation: Destroy temporary directory [tmp_test]
    [***] Success

--> Results:
    Groups: ** / **
    [***]   Successful: 3   Errors: 0   Commands: 3/3

Steps to reproduce the behavior (include code & usage example):

Expected behavior

I would expect the order of operations remain unchanged whether I put them in python functions or not.

Meta

• Include output of pyinfra --support.

--> Support information:

    If you are having issues with pyinfra or wish to make feature requests, please
    check out the GitHub issues at https://github.com/Fizzadar/pyinfra/issues .
    When adding an issue, be sure to include the following:

    System: Linux
      Platform: Linux-3.10.0-1160.59.1.el7.x86_64-x86_64-with-glibc2.17
      Release: 3.10.0-1160.59.1.el7.x86_64
      Machine: x86_64
    pyinfra: v2.0.1
    Executable: /home/**/.local/bin/pyinfra
    Python: 3.9.4 (CPython, GCC 4.8.5 20150623 (Red Hat 4.8.5-44))

• How was pyinfra installed (source/pip)? pip

[Fizzadar]

This should now be fixed in v2.2!

[julienlavergne]

ok