diff --git a/compute_endpoint/tests/unit/test_gce_container.py b/compute_endpoint/tests/unit/test_gce_container.py index 82394f710..23dc45bd1 100644 --- a/compute_endpoint/tests/unit/test_gce_container.py +++ b/compute_endpoint/tests/unit/test_gce_container.py @@ -15,6 +15,8 @@ @pytest.fixture def gce_factory(tmp_path, randomstring) -> t.Callable: + engines: list[GlobusComputeEngine] = [] + def _kernel(**k): expect_uri = randomstring(length=random.randint(1, 20)) expect_opts = randomstring(length=random.randint(1, 20)) @@ -27,16 +29,21 @@ def _kernel(**k): **k, } with mock.patch(f"{_MOCK_BASE}ReportingThread"): - gce = GlobusComputeEngine(**k) - with mock.patch.object(gce.executor, "start"): - gce.start(endpoint_id=uuid.uuid4(), run_dir=tmp_path) - assert gce.executor.start.called + with mock.patch(f"{_MOCK_BASE}JobStatusPoller"): + gce = GlobusComputeEngine(**k) + gce.executor.start = mock.Mock() + gce.start(endpoint_id=uuid.uuid4(), run_dir=str(tmp_path)) + assert gce.executor.start.called + + engines.append(gce) - # No cleanup necessary because executor not started return gce, expect_uri, expect_opts yield _kernel + for e in engines: + e.shutdown() + def test_docker(tmp_path, gce_factory): gce, exp_uri, exp_opts = gce_factory(container_type="docker")