diff --git a/arkindex/process/api.py b/arkindex/process/api.py
index 8348ff930641d50335a665336eb1f0b95502208f..b1dfb348b851fa2ca1a69f39f43a48305a137ab7 100644
--- a/arkindex/process/api.py
+++ b/arkindex/process/api.py
@@ -922,7 +922,7 @@ class WorkerTypesList(ListAPIView):
"mode",
type=str,
default="complete",
- description=("If set to `simple`, exclude versions with no tag")
+ description="When this is set to `simple`, only the worker versions with a tag set, or on the `main` or `master` branches, will be returned."
)
],
),
@@ -967,8 +967,8 @@ class WorkerVersionList(ListCreateAPIView):
def get_queryset(self):
filters = Q()
if self.simple_mode:
- # Limit output to versions with tags
- filters = Q(tag__isnull=False)
+ # Limit output to versions with tags or from the main/master branch
+ filters = Q(tag__isnull=False) | Q(branch__in=["main", "master"])
return self.worker.versions \
.using("default") \
.filter(filters) \
diff --git a/arkindex/process/tests/worker_versions/test_list.py b/arkindex/process/tests/worker_versions/test_list.py
index e670faefdcf44893aad992bec24fd3fddf05998a..4fc3d5800b9d252f60aec39e4a43a3dedf6a1f02 100644
--- a/arkindex/process/tests/worker_versions/test_list.py
+++ b/arkindex/process/tests/worker_versions/test_list.py
@@ -19,6 +19,16 @@ class TestWorkerVersionList(FixtureAPITestCase):
cls.version_1 = cls.worker_reco.versions.get()
cls.version_2 = cls.worker_dla.versions.get()
+ cls.version_main_branch = cls.worker_reco.versions.create(
+ configuration={},
+ branch="main",
+ revision_url="https://gitlab.com/NERV/eva/commit/12"
+ )
+ cls.version_other_branch = cls.worker_reco.versions.create(
+ configuration={},
+ branch="eva-00",
+ revision_url="https://gitlab.com/NERV/eva/commit/000002"
+ )
farm = Farm.objects.first()
process = cls.corpus.processes.create(
@@ -52,7 +62,7 @@ class TestWorkerVersionList(FixtureAPITestCase):
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertDictEqual(response.json(), {
- "count": 2,
+ "count": 4,
"next": None,
"number": 1,
"previous": None,
@@ -79,6 +89,50 @@ class TestWorkerVersionList(FixtureAPITestCase):
"revision_url": "https://gitlab.com/NERV/eva/commit/63e377e7f88c743d8428fc4e4eaedfc1c9356754",
"created": "2050-09-09T09:09:09.090909Z",
},
+ {
+ "id": str(self.version_other_branch.id),
+ "configuration": {},
+ "docker_image_iid": None,
+ "state": "created",
+ "gpu_usage": FeatureUsage.Disabled.value,
+ "model_usage": FeatureUsage.Disabled.value,
+ "worker": {
+ "id": str(self.worker_reco.id),
+ "name": self.worker_reco.name,
+ "type": self.worker_reco.type.slug,
+ "slug": self.worker_reco.slug,
+ "description": self.worker_reco.description,
+ "archived": bool(self.worker_reco.archived),
+ "repository_url": self.worker_reco.repository_url,
+ },
+ "version": None,
+ "tag": None,
+ "branch": "eva-00",
+ "revision_url": "https://gitlab.com/NERV/eva/commit/000002",
+ "created": self.version_other_branch.created.isoformat().replace("+00:00", "Z"),
+ },
+ {
+ "id": str(self.version_main_branch.id),
+ "configuration": {},
+ "docker_image_iid": None,
+ "state": "created",
+ "gpu_usage": FeatureUsage.Disabled.value,
+ "model_usage": FeatureUsage.Disabled.value,
+ "worker": {
+ "id": str(self.worker_reco.id),
+ "name": self.worker_reco.name,
+ "type": self.worker_reco.type.slug,
+ "slug": self.worker_reco.slug,
+ "description": self.worker_reco.description,
+ "archived": bool(self.worker_reco.archived),
+ "repository_url": self.worker_reco.repository_url,
+ },
+ "version": None,
+ "tag": None,
+ "branch": "main",
+ "revision_url": "https://gitlab.com/NERV/eva/commit/12",
+ "created": self.version_main_branch.created.isoformat().replace("+00:00", "Z"),
+ },
{
"id": str(self.version_1.id),
"configuration": {"test": 42},
@@ -139,8 +193,75 @@ class TestWorkerVersionList(FixtureAPITestCase):
self.assertEqual(response.status_code, status.HTTP_200_OK)
data = response.json()
- self.assertEqual(data["count"], 1)
- self.assertEqual(data["results"][0]["id"], str(self.worker_reco.versions.first().id))
+ self.assertEqual(data["count"], 3)
+ self.assertListEqual(data["results"], [
+ {
+ "id": str(self.version_other_branch.id),
+ "configuration": {},
+ "docker_image_iid": None,
+ "state": "created",
+ "gpu_usage": FeatureUsage.Disabled.value,
+ "model_usage": FeatureUsage.Disabled.value,
+ "worker": {
+ "id": str(self.worker_reco.id),
+ "name": self.worker_reco.name,
+ "type": self.worker_reco.type.slug,
+ "slug": self.worker_reco.slug,
+ "description": self.worker_reco.description,
+ "archived": bool(self.worker_reco.archived),
+ "repository_url": self.worker_reco.repository_url,
+ },
+ "version": None,
+ "tag": None,
+ "branch": "eva-00",
+ "revision_url": "https://gitlab.com/NERV/eva/commit/000002",
+ "created": self.version_other_branch.created.isoformat().replace("+00:00", "Z"),
+ },
+ {
+ "id": str(self.version_main_branch.id),
+ "configuration": {},
+ "docker_image_iid": None,
+ "state": "created",
+ "gpu_usage": FeatureUsage.Disabled.value,
+ "model_usage": FeatureUsage.Disabled.value,
+ "worker": {
+ "id": str(self.worker_reco.id),
+ "name": self.worker_reco.name,
+ "type": self.worker_reco.type.slug,
+ "slug": self.worker_reco.slug,
+ "description": self.worker_reco.description,
+ "archived": bool(self.worker_reco.archived),
+ "repository_url": self.worker_reco.repository_url,
+ },
+ "version": None,
+ "tag": None,
+ "branch": "main",
+ "revision_url": "https://gitlab.com/NERV/eva/commit/12",
+ "created": self.version_main_branch.created.isoformat().replace("+00:00", "Z"),
+ },
+ {
+ "id": str(self.version_1.id),
+ "configuration": {"test": 42},
+ "docker_image_iid": self.version_1.docker_image_iid,
+ "state": "available",
+ "gpu_usage": FeatureUsage.Disabled.value,
+ "model_usage": FeatureUsage.Disabled.value,
+ "worker": {
+ "id": str(self.worker_reco.id),
+ "name": self.worker_reco.name,
+ "type": self.worker_reco.type.slug,
+ "slug": self.worker_reco.slug,
+ "description": self.worker_reco.description,
+ "archived": bool(self.worker_reco.archived),
+ "repository_url": self.worker_reco.repository_url,
+ },
+ "version": 1,
+ "tag": None,
+ "branch": None,
+ "revision_url": None,
+ "created": self.version_1.created.isoformat().replace("+00:00", "Z"),
+ }
+ ])
def test_list_filter_worker(self):
self.client.force_login(self.user)
@@ -167,7 +288,8 @@ class TestWorkerVersionList(FixtureAPITestCase):
def test_list_simple_mode(self):
"""
- With the mode attribute set to simple, worker versions with no tag are excluded
+ With the mode attribute set to simple, worker versions with no tag or not from the
+ master / main branch are excluded
"""
WorkerVersion.objects.bulk_create([
WorkerVersion(
@@ -186,10 +308,14 @@ class TestWorkerVersionList(FixtureAPITestCase):
)
self.assertEqual(response.status_code, status.HTTP_200_OK)
data = response.json()
- self.assertEqual(data["count"], 4)
+ self.assertEqual(data["count"], 6)
self.assertCountEqual(
[version["tag"] for version in data["results"]],
- ["v0.1", "test", None, None],
+ ["v0.1", "test", None, None, None, None],
+ )
+ self.assertCountEqual(
+ [version["branch"] for version in data["results"]],
+ [None, None, None, None, "main", "eva-00"],
)
# Simple mode filters out versions without a tag
@@ -201,8 +327,12 @@ class TestWorkerVersionList(FixtureAPITestCase):
)
self.assertEqual(response.status_code, status.HTTP_200_OK)
data = response.json()
- self.assertEqual(data["count"], 2)
+ self.assertEqual(data["count"], 3)
self.assertCountEqual(
[version["tag"] for version in data["results"]],
- ["v0.1", "test"],
+ ["v0.1", "test", None],
+ )
+ self.assertCountEqual(
+ [version["branch"] for version in data["results"]],
+ [None, None, "main"],
)