From 2a48d2f7e8067c9d2c6ce45dc40d35293201b06d Mon Sep 17 00:00:00 2001
From: Alejandro Aristizabal <alejandro.aristizabal24@gmail.com>
Date: Fri, 28 Jul 2023 16:12:44 -0500
Subject: [PATCH] Mock pexpect context management. Fix tests

---
 cli/medperf/tests/entities/test_cube.py | 2 +-
 cli/medperf/tests/mocks/pexpect.py      | 6 ++++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/cli/medperf/tests/entities/test_cube.py b/cli/medperf/tests/entities/test_cube.py
index 37be80bd2..a2efac669 100644
--- a/cli/medperf/tests/entities/test_cube.py
+++ b/cli/medperf/tests/entities/test_cube.py
@@ -133,7 +133,7 @@ def test_cube_runs_command_with_extra_args(self, mocker, setup, task):
         spy = mocker.patch("pexpect.spawn", side_effect=mpexpect.spawn)
         expected_cmd = (
             f"mlcube run --mlcube={self.manifest_path} --task={task} "
-            + f'--platform={self.platform} test="test"'
+            + f'--platform={self.platform} --network=none test="test"'
         )
 
         # Act
diff --git a/cli/medperf/tests/mocks/pexpect.py b/cli/medperf/tests/mocks/pexpect.py
index 23e946099..f96fb0dc6 100644
--- a/cli/medperf/tests/mocks/pexpect.py
+++ b/cli/medperf/tests/mocks/pexpect.py
@@ -2,6 +2,12 @@ class MockChild:
     def __init__(self, exitstatus):
         self.exitstatus = exitstatus
 
+    def __enter__(self):
+        pass
+
+    def __exit__(self):
+        self.close()
+
     def isalive(self):
         return False