Skip to content

Commit

Permalink
fix test
Browse files Browse the repository at this point in the history
  • Loading branch information
epwalsh committed Jan 9, 2025
1 parent 919e913 commit 950c644
Show file tree
Hide file tree
Showing 2 changed files with 1 addition and 28 deletions.
25 changes: 0 additions & 25 deletions conftest.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import logging
import random
import sys

import pytest
Expand Down Expand Up @@ -35,35 +34,11 @@ def beaker(workspace_name):
return beaker_client


def _get_free_cluster(beaker: Beaker) -> str:
choices = [
"ai2/general-cirrascale",
"ai2/allennlp-cirrascale",
"ai2/aristo-cirrascale",
"ai2/mosaic-cirrascale",
"ai2/s2-cirrascale",
]
random.shuffle(choices)
for cluster in choices:
utilization = beaker.cluster.utilization(cluster)
if utilization.queued_jobs == 0:
logger.info("Found suitable on-prem cluster '%s'", cluster)
return cluster
return "ai2/general-cirrascale"


@pytest.fixture()
def beaker_cluster_name(beaker: Beaker) -> str:
return _get_free_cluster(beaker)


if __name__ == "__main__":
beaker_client = Beaker.from_env()
assert len(sys.argv) == 2
fixture = sys.argv[-1]
if fixture == "run_name":
print(_get_unique_name())
elif fixture == "cluster":
print(_get_free_cluster(beaker_client))
else:
raise ValueError(f"Bad fixture argument '{fixture}'")
4 changes: 1 addition & 3 deletions tests/main_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ def test_version():
assert VERSION in result.stdout


def test_dry_run(workspace_name: str, beaker_cluster_name: str, run_name: str):
def test_dry_run(workspace_name: str, run_name: str):
result = subprocess.run(
[
"gantry",
Expand All @@ -26,8 +26,6 @@ def test_dry_run(workspace_name: str, beaker_cluster_name: str, run_name: str):
run_name,
"--workspace",
workspace_name,
"--cluster",
beaker_cluster_name,
"--yes",
"--",
"python",
Expand Down

0 comments on commit 950c644

Please sign in to comment.