Update tests to add ConnectionOpts argument.

facebookresearch · May 12, 2022 · 8735fce · 8735fce
1 parent 31a9430
commit 8735fce
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/tests/compiler_env_test.py b/tests/compiler_env_test.py
@@ -9,7 +9,7 @@
 
 from compiler_gym.envs import llvm
 from compiler_gym.envs.llvm import LlvmEnv
-from compiler_gym.service.connection import CompilerGymServiceConnection
+from compiler_gym.service.connection import CompilerGymServiceConnection, ConnectionOpts
 from tests.test_main import main
 
 pytest_plugins = ["tests.pytest_plugins.llvm"]
@@ -174,7 +174,7 @@ def test_step_session_id_not_found(env: LlvmEnv):
 @pytest.fixture(scope="function")
 def remote_env() -> LlvmEnv:
     """A test fixture that yields a connection to a remote service."""
-    service = CompilerGymServiceConnection(llvm.LLVM_SERVICE_BINARY)
+    service = CompilerGymServiceConnection(llvm.LLVM_SERVICE_BINARY, ConnectionOpts())
     try:
         with LlvmEnv(service=service.connection.url) as env:
             yield env

diff --git a/tests/mlir/mlir_env_test.py b/tests/mlir/mlir_env_test.py
@@ -12,7 +12,7 @@
 import compiler_gym
 from compiler_gym.envs import CompilerEnv, mlir
 from compiler_gym.envs.mlir import MlirEnv
-from compiler_gym.service.connection import CompilerGymServiceConnection
+from compiler_gym.service.connection import CompilerGymServiceConnection, ConnectionOpts
 from compiler_gym.spaces import (
     Box,
     Dict,
@@ -36,7 +36,9 @@ def env(request) -> CompilerEnv:
         with gym.make("mlir-v0") as env:
             yield env
     else:
-        service = CompilerGymServiceConnection(mlir.MLIR_SERVICE_BINARY)
+        service = CompilerGymServiceConnection(
+            mlir.MLIR_SERVICE_BINARY, ConnectionOpts()
+        )
         try:
             with MlirEnv(service=service.connection.url) as env:
                 yield env