| 114 | |
| 115 | |
| 116 | class HqEnv(Env): |
| 117 | def __init__(self, work_dir, mock: ProgramMock, debug=True): |
| 118 | Env.__init__(self, work_dir) |
| 119 | self.mock = mock |
| 120 | self.server = None |
| 121 | self.id_counter = 0 |
| 122 | self.do_final_check = True |
| 123 | self.server_dir = "" |
| 124 | self.debug = debug |
| 125 | |
| 126 | def no_final_check(self): |
| 127 | self.do_final_check = False |
| 128 | |
| 129 | def make_default_env(self, log=True): |
| 130 | env = os.environ.copy() |
| 131 | if log: |
| 132 | env["RUST_LOG"] = "tako=trace,hyperqueue=trace" |
| 133 | env["RUST_BACKTRACE"] = "full" |
| 134 | env["HQ_TEST"] = "1" |
| 135 | self.mock.update_env(env) |
| 136 | return env |
| 137 | |
| 138 | @staticmethod |
| 139 | def server_args(server_dir="hq-server", debug=True): |
| 140 | args = [ |
| 141 | get_hq_binary(debug=debug), |
| 142 | "--colors", |
| 143 | "never", |
| 144 | "--server-dir", |
| 145 | server_dir, |
| 146 | ] |
| 147 | if debug: |
| 148 | args.append("--debug") |
| 149 | |
| 150 | args += ["server", "start"] |
| 151 | return args |
| 152 | |
| 153 | def start_server(self, server_dir="hq-server", args=None, env=None) -> subprocess.Popen: |
| 154 | self.server_dir = os.path.join(self.work_path, server_dir) |
| 155 | environment = self.make_default_env() |
| 156 | if env: |
| 157 | environment.update(env) |
| 158 | server_args = self.server_args(self.server_dir, debug=self.debug) |
| 159 | if args: |
| 160 | server_args += args |
| 161 | process = self.start_process("server", server_args, env=environment) |
| 162 | time.sleep(0.2) |
| 163 | self.check_running_processes() |
| 164 | return process |
| 165 | |
| 166 | def start_workers(self, count, **kwargs) -> List[subprocess.Popen]: |
| 167 | workers = [] |
| 168 | for _ in range(count): |
| 169 | workers.append(self.start_worker(**kwargs)) |
| 170 | return workers |
| 171 | |
| 172 | def start_worker( |
| 173 | self, |
no outgoing calls