tests: use, rather than set, context

Author: Fizzadar

tests/test_api/test_api_operations.py

@@ -426,32 +426,31 @@ def test_nested_op_api(self):
 
         somehost = inventory.get_host("somehost")
 
-        ctx_state.set(state)
-        ctx_host.set(somehost)
-
         pyinfra.is_cli = True
 
-        try:
-            outer_result = server.shell(commands="echo outer")
-            assert outer_result._combined_output is None
+        with ctx_state.use(state):
+            with ctx_host.use(somehost):
+                try:
+                    outer_result = server.shell(commands="echo outer")
+                    assert outer_result._combined_output is None
 
-            def callback():
-                inner_result = server.shell(commands="echo inner")
-                assert inner_result._combined_output is not None
+                    def callback():
+                        inner_result = server.shell(commands="echo inner")
+                        assert inner_result._combined_output is not None
 
-            python.call(function=callback)
+                    python.call(function=callback)
 
-            assert len(state.get_op_order()) == 2
+                    assert len(state.get_op_order()) == 2
 
-            run_ops(state)
+                    run_ops(state)
 
-            assert len(state.get_op_order()) == 3
-            assert state.results[somehost].success_ops == 3
-            assert outer_result._combined_output is not None
+                    assert len(state.get_op_order()) == 3
+                    assert state.results[somehost].success_ops == 3
+                    assert outer_result._combined_output is not None
 
-            disconnect_all(state)
-        finally:
-            pyinfra.is_cli = False
+                    disconnect_all(state)
+                finally:
+                    pyinfra.is_cli = False
 
 
 class TestOperationFailures(PatchSSHTestCase):
@@ -519,40 +518,40 @@ def test_cli_op_line_numbers(self):
         state.current_deploy_filename = __file__
 
         pyinfra.is_cli = True
-        ctx_state.set(state)
 
-        # Add op to both hosts
-        for name in ("anotherhost", "somehost"):
-            ctx_host.set(inventory.get_host(name))
-            server.shell("echo hi")  # note this is called twice but on *the same line*
+        with ctx_state.use(state):
+            # Add op to both hosts
+            for name in ("anotherhost", "somehost"):
+                with ctx_host.use(inventory.get_host(name)):
+                    server.shell("echo hi")  # note this is called twice but on *the same line*
 
-        # Add op to just the second host - using the context modules such that
-        # it replicates a deploy file.
-        ctx_host.set(inventory.get_host("anotherhost"))
-        first_context_hash = server.user("anotherhost_user")._hash
+            # Add op to just the second host - using the context modules such that
+            # it replicates a deploy file.
+            ctx_host.set(inventory.get_host("anotherhost"))
+            first_context_hash = server.user("anotherhost_user")._hash
 
-        # Add op to just the first host - using the context modules such that
-        # it replicates a deploy file.
-        ctx_host.set(inventory.get_host("somehost"))
-        second_context_hash = server.user("somehost_user")._hash
+            # Add op to just the first host - using the context modules such that
+            # it replicates a deploy file.
+            ctx_host.set(inventory.get_host("somehost"))
+            second_context_hash = server.user("somehost_user")._hash
 
-        ctx_state.reset()
-        ctx_host.reset()
+            ctx_state.reset()
+            ctx_host.reset()
 
-        pyinfra.is_cli = False
+            pyinfra.is_cli = False
 
-        print(state.ops)
-        # Ensure there are two ops
-        op_order = state.get_op_order()
-        assert len(op_order) == 3
+            print(state.ops)
+            # Ensure there are two ops
+            op_order = state.get_op_order()
+            assert len(op_order) == 3
 
-        # And that the two ops above were called in the expected order
-        assert op_order[1] == first_context_hash
-        assert op_order[2] == second_context_hash
+            # And that the two ops above were called in the expected order
+            assert op_order[1] == first_context_hash
+            assert op_order[2] == second_context_hash
 
-        # Ensure somehost has two ops and anotherhost only has the one
-        assert len(state.ops[inventory.get_host("somehost")]) == 2
-        assert len(state.ops[inventory.get_host("anotherhost")]) == 2
+            # Ensure somehost has two ops and anotherhost only has the one
+            assert len(state.ops[inventory.get_host("somehost")]) == 2
+            assert len(state.ops[inventory.get_host("anotherhost")]) == 2
 
     # In API mode, pyinfra *overrides* the line numbers such that whenever an
     # operation or deploy is added it is simply appended. This makes sense as

tests/test_cli/test_context_objects.py

@@ -5,8 +5,8 @@
 from pyinfra.context import ctx_host, ctx_inventory
 
 
-def _create_host():
-    return Host(None, None, None, None)
+def _create_host(name: str = "host"):
+    return Host(name, Inventory(([name], {})), None, None)
 
 
 class TestHostContextObject(TestCase):
@@ -15,59 +15,59 @@ def test_context_host_dir_matches_class(self):
 
     def test_context_host_eq(self):
         host_obj = _create_host()
-        ctx_host.set(host_obj)
-        assert host == host_obj
+        with ctx_host.use(host_obj):
+            assert host == host_obj
 
     def test_context_host_repr(self):
-        host_obj = _create_host()
-        ctx_host.set(host_obj)
-        assert repr(host) == "ContextObject(Host):Host(None)"
+        host_obj = _create_host("somehost")
+        with ctx_host.use(host_obj):
+            assert repr(host) == "ContextObject(Host):Host(somehost)"
 
     def test_context_host_str(self):
         host_obj = _create_host()
-        ctx_host.set(host_obj)
-        assert str(host_obj) == "None"
+        with ctx_host.use(host_obj):
+            assert str(host_obj) == "host"
 
     def test_context_host_attr(self):
         host_obj = _create_host()
-        ctx_host.set(host_obj)
 
         with self.assertRaises(AttributeError):
             host_obj.hello
 
-        setattr(host, "hello", "world")
-        assert host_obj.hello == host.hello
+        with ctx_host.use(host_obj):
+            setattr(host, "hello", "world")
+            assert host_obj.hello == host.hello
 
     def test_context_host_class_attr(self):
         host_obj = _create_host()
-        ctx_host.set(host_obj)
-        assert ctx_host.isset() is True
 
-        with self.assertRaises(AttributeError):
-            host_obj.hello
+        with ctx_host.use(host_obj):
+            assert ctx_host.isset() is True
+
+            with self.assertRaises(AttributeError):
+                host_obj.hello
 
-        setattr(Host, "hello", "class_world")
-        setattr(host_obj, "hello", "instance_world")
+            setattr(Host, "hello", "class_world")
+            setattr(host_obj, "hello", "instance_world")
 
-        assert host.hello == host.hello
+            assert host.hello == host.hello
 
         # Reset and check fallback to class variable
-        ctx_host.reset()
         assert ctx_host.isset() is False
         assert host.hello == "class_world"
 
 
 class TestInventoryContextObject(TestCase):
     def test_context_inventory_len(self):
         inventory_obj = Inventory(("host", "anotherhost"))
-        ctx_inventory.set(inventory_obj)
-        assert ctx_inventory.isset() is True
 
-        assert len(inventory) == len(inventory_obj)
+        with ctx_inventory.use(inventory_obj):
+            assert ctx_inventory.isset() is True
+            assert len(inventory) == len(inventory_obj)
 
     def test_context_inventory_iter(self):
         inventory_obj = Inventory(("host", "anotherhost"))
-        ctx_inventory.set(inventory_obj)
-        assert ctx_inventory.isset() is True
 
-        assert list(iter(inventory)) == list(iter(inventory_obj))
+        with ctx_inventory.use(inventory_obj):
+            assert ctx_inventory.isset() is True
+            assert list(iter(inventory)) == list(iter(inventory_obj))