def test_session_create(client):
user = UserFactory(is_staff=True)
ws = WorkstationFactory()
# Create some workstations and pretend that they're ready
wsi_old = WorkstationImageFactory(workstation=ws, ready=True)
wsi_new = WorkstationImageFactory(workstation=ws, ready=True)
wsi_not_ready = WorkstationImageFactory(workstation=ws)
wsi_other_ws = WorkstationImageFactory(ready=True)
assert Session.objects.count() == 0
response = get_view_for_user(
client=client,
method=client.post,
viewname="workstations:session-create",
reverse_kwargs={"slug": ws.slug},
user=user,
)
assert response.status_code == 302
sessions = Session.objects.all()
assert len(sessions) == 1
# Should select the most recent workstation
assert sessions[0].workstation_image == wsi_new
assert sessions[0].creator == user
def test_workstationimage_update(client):
user = UserFactory(is_staff=True)
wsi = WorkstationImageFactory()
assert wsi.initial_path != ""
assert wsi.websocket_port != 1337
assert wsi.http_port != 1234
response = get_view_for_user(
client=client,
method=client.post,
viewname="workstations:image-update",
reverse_kwargs={"slug": wsi.workstation.slug, "pk": wsi.pk},
user=user,
data={"initial_path": "", "websocket_port": 1337, "http_port": 1234},
)
assert response.status_code == 302
assert response.url == wsi.get_absolute_url()
wsi.refresh_from_db()
assert wsi.initial_path == ""
assert wsi.websocket_port == 1337
assert wsi.http_port == 1234
def test_workstationimage_update(client):
user = UserFactory(is_staff=True)
wsi = WorkstationImageFactory()
assert wsi.initial_path != ""
assert wsi.websocket_port != 1337
assert wsi.http_port != 1234
response = get_view_for_user(
client=client,
method=client.post,
viewname="workstations:image-update",
reverse_kwargs={"slug": wsi.workstation.slug, "pk": wsi.pk},
user=user,
data={"initial_path": "", "websocket_port": 1337, "http_port": 1234},
)
assert response.status_code == 302
assert response.url == wsi.get_absolute_url()
wsi.refresh_from_db()
assert wsi.initial_path == ""
assert wsi.websocket_port == 1337
assert wsi.http_port == 1234
def test_session_redirect_staff_only(client):
# Create the default image
WorkstationImageFactory(
workstation__title=settings.DEFAULT_WORKSTATION_SLUG, ready=True
)
wsi = WorkstationImageFactory(ready=True)
# Validate
validate_staff_only_view(
client=client,
viewname="workstations:default-session-redirect",
reverse_kwargs={},
should_redirect=True,
)
validate_staff_only_view(
client=client,
viewname="workstations:workstation-session-redirect",
reverse_kwargs={"slug": wsi.workstation.slug},
should_redirect=True,
)
validate_staff_only_view(
client=client,
viewname="workstations:workstation-image-session-redirect",
reverse_kwargs={"slug": wsi.workstation.slug, "pk": wsi.pk},
should_redirect=True,
)
def test_session_create(client):
user = UserFactory()
ws = WorkstationFactory()
ws.add_user(user=user)
# Create some workstations and pretend that they're ready
WorkstationImageFactory(workstation=ws, ready=True) # Old WSI
wsi_new = WorkstationImageFactory(workstation=ws, ready=True)
WorkstationImageFactory(workstation=ws) # WSI not ready
WorkstationImageFactory(ready=True) # Image for some other ws
assert Session.objects.count() == 0
response = get_view_for_user(
client=client,
method=client.post,
viewname="workstations:workstation-session-create",
reverse_kwargs={"slug": ws.slug},
user=user,
data={"region": "eu-central-1"},
)
assert response.status_code == 302
sessions = Session.objects.all()
assert len(sessions) == 1
# Should select the most recent workstation
assert sessions[0].workstation_image == wsi_new
assert sessions[0].creator == user
def test_get_workstation_image_or_404():
# No default workstation
with pytest.raises(Http404):
get_workstation_image_or_404()
default_workstation = Workstation.objects.get(
slug=settings.DEFAULT_WORKSTATION_SLUG
)
default_wsi = WorkstationImageFactory(
workstation=default_workstation, ready=True
)
wsi = WorkstationImageFactory(ready=True)
found_wsi = get_workstation_image_or_404()
assert found_wsi == default_wsi
found_wsi = get_workstation_image_or_404(slug=wsi.workstation.slug)
assert found_wsi == wsi
found_wsi = get_workstation_image_or_404(pk=wsi.pk)
assert found_wsi == wsi
# No images for workstation
with pytest.raises(Http404):
get_workstation_image_or_404(slug=WorkstationFactory().slug)
# Incorrect pk
with pytest.raises(Http404):
get_workstation_image_or_404(pk=WorkstationFactory().pk)
def test_get_or_create_active_session():
user = UserFactory()
wsi = WorkstationImageFactory()
assert Session.objects.all().count() == 0
s = get_or_create_active_session(
user=user, workstation_image=wsi, region="eu-central-1"
)
assert s.workstation_image == wsi
assert s.creator == user
assert Session.objects.all().count() == 1
# Same workstation image and user
s_1 = get_or_create_active_session(
user=user, workstation_image=wsi, region="eu-central-1"
)
assert s == s_1
# Different workstation image, same user
wsi_1 = WorkstationImageFactory()
s_2 = get_or_create_active_session(
user=user, workstation_image=wsi_1, region="eu-central-1"
)
assert s_2.workstation_image == wsi_1
assert s_2.creator == user
assert Session.objects.all().count() == 2
assert s_1 != s_2
# Same workstation image, different user
user_1 = UserFactory()
s_3 = get_or_create_active_session(
user=user_1, workstation_image=wsi, region="eu-central-1"
)
assert s_3.workstation_image == wsi
assert s_3.creator == user_1
assert Session.objects.all().count() == 3
# Stop the original session, original workstation image and user
s.status = s.STOPPED
s.save()
s_4 = get_or_create_active_session(
user=user, workstation_image=wsi, region="eu-central-1"
)
assert s_4.workstation_image == wsi
assert s_4.creator == user
assert Session.objects.all().count() == 4
def test_correct_session_stopped(http_image, docker_client, settings):
path, sha256 = http_image
wsi = WorkstationImageFactory(image__from_path=path,
image_sha256=sha256,
ready=True)
# Execute the celery in place
settings.task_eager_propagates = (True, )
settings.task_always_eager = (True, )
try:
s1, s2 = (
SessionFactory(workstation_image=wsi),
SessionFactory(workstation_image=wsi),
)
assert s1.service.container
assert s2.service.container
s2.user_finished = True
s2.save()
assert s1.service.container
with pytest.raises(NotFound):
# noinspection PyStatementEffect
s2.service.container
finally:
stop_all_sessions()
def test_session_cleanup(http_image, docker_client, settings):
path, sha256 = http_image
wsi = WorkstationImageFactory(image__from_path=path,
image_sha256=sha256,
ready=True)
# Execute the celery in place
settings.task_eager_propagates = (True, )
settings.task_always_eager = (True, )
try:
s1, s2 = (
SessionFactory(workstation_image=wsi),
SessionFactory(workstation_image=wsi,
maximum_duration=timedelta(seconds=0)),
)
assert s1.service.container
assert s2.service.container
stop_expired_services(app_label="workstations", model_name="session")
assert s1.service.container
with pytest.raises(NotFound):
# noinspection PyStatementEffect
s2.service.container
finally:
stop_all_sessions()
def test_session_limit(http_image, docker_client, settings):
path, sha256 = http_image
wsi = WorkstationImageFactory(image__from_path=path,
image_sha256=sha256,
ready=True)
# Execute the celery in place
settings.WORKSTATIONS_MAXIMUM_SESSIONS = 1
settings.task_eager_propagates = (True, )
settings.task_always_eager = (True, )
try:
with capture_on_commit_callbacks(execute=True):
s1 = SessionFactory(workstation_image=wsi)
s1.refresh_from_db()
assert s1.status == s1.STARTED
with capture_on_commit_callbacks(execute=True):
s2 = SessionFactory(workstation_image=wsi)
s2.refresh_from_db()
assert s2.status == s2.FAILED
s1.stop()
with capture_on_commit_callbacks(execute=True):
s3 = SessionFactory(workstation_image=wsi)
s3.refresh_from_db()
assert s3.status == s3.STARTED
finally:
stop_all_sessions()
def test_workstations_staff_views(client, view):
if view in [
"workstations:update",
"workstations:detail",
"workstations:image-create",
"workstations:session-create",
]:
reverse_kwargs = {"slug": WorkstationFactory().slug}
elif view in ["workstations:image-detail", "workstations:image-update"]:
wsi = WorkstationImageFactory()
reverse_kwargs = {"slug": wsi.workstation.slug, "pk": wsi.pk}
elif view in [
"workstations:session-detail",
"workstations:session-update",
]:
session = SessionFactory()
reverse_kwargs = {
"slug": session.workstation_image.workstation.slug,
"pk": session.pk,
}
else:
reverse_kwargs = {}
validate_staff_only_view(
client=client, viewname=view, reverse_kwargs=reverse_kwargs
)
def workstation_set():
ws = WorkstationFactory()
wsi = WorkstationImageFactory(workstation=ws)
e, u, u1 = UserFactory(), UserFactory(), UserFactory()
wss = WorkstationSet(workstation=ws, editor=e, user=u, user1=u1, image=wsi)
wss.workstation.add_editor(user=e)
wss.workstation.add_user(user=u)
return wss
def test_workstationimage_detail(client):
user = UserFactory(is_staff=True)
ws = WorkstationFactory()
wsi1, wsi2 = (
WorkstationImageFactory(workstation=ws),
WorkstationImageFactory(workstation=ws),
)
response = get_view_for_user(
viewname="workstations:image-detail",
reverse_kwargs={"slug": ws.slug, "pk": wsi1.pk},
client=client,
user=user,
)
assert response.status_code == 200
assert str(wsi1.pk) in response.rendered_content
assert str(wsi2.pk) not in response.rendered_content
def workstation_set():
ws = WorkstationFactory()
wsi = WorkstationImageFactory(workstation=ws)
e, u, u1 = UserFactory(), UserFactory(), UserFactory()
for user in [e, u, u1]:
VerificationFactory(user=user, is_verified=True)
wss = WorkstationSet(workstation=ws, editor=e, user=u, user1=u1, image=wsi)
wss.workstation.add_editor(user=e)
wss.workstation.add_user(user=u)
return wss
def test_session_cleanup(http_image, docker_client, settings):
path, sha256 = http_image
wsi = WorkstationImageFactory(image__from_path=path,
image_sha256=sha256,
ready=True)
# Execute the celery in place
settings.task_eager_propagates = (True, )
settings.task_always_eager = (True, )
default_region = "eu-nl-1"
try:
with capture_on_commit_callbacks(execute=True):
s1, s2, s3 = (
SessionFactory(workstation_image=wsi, region=default_region),
SessionFactory(
workstation_image=wsi,
maximum_duration=timedelta(seconds=0),
region=default_region,
),
# An expired service in a different region
SessionFactory(
workstation_image=wsi,
maximum_duration=timedelta(seconds=0),
region="us-east-1",
),
)
assert s1.service.container
assert s2.service.container
assert s3.service.container
# Stop expired services in the default region
stop_expired_services(
app_label="workstations",
model_name="session",
region=default_region,
)
assert s1.service.container
with pytest.raises(NotFound):
# noinspection PyStatementEffect
s2.service.container
assert s3.service.container
finally:
stop_all_sessions()
def test_workstation_ready(http_image, docker_client, settings):
path, sha256 = http_image
wsi = WorkstationImageFactory(image__from_path=path,
image_sha256=sha256,
ready=False)
# Execute the celery in place
settings.task_eager_propagates = (True, )
settings.task_always_eager = (True, )
s = SessionFactory(workstation_image=wsi)
s.refresh_from_db()
assert s.status == s.FAILED
def test_session_redirect(client):
user = UserFactory(is_staff=True)
WorkstationImageFactory(
workstation__title=settings.DEFAULT_WORKSTATION_SLUG, ready=True)
response = get_view_for_user(
client=client,
viewname="workstations:default-session-redirect",
user=user,
)
assert response.status_code == 302
response = get_view_for_user(client=client, user=user, url=response.url)
assert response.status_code == 200
def test_session_start(http_image, docker_client, settings):
path, sha256 = http_image
wsi = WorkstationImageFactory(image__from_path=path,
image_sha256=sha256,
ready=True)
# Execute the celery in place
settings.task_eager_propagates = (True, )
settings.task_always_eager = (True, )
with capture_on_commit_callbacks(execute=True):
s = SessionFactory(workstation_image=wsi)
try:
assert s.service.container
s.refresh_from_db()
assert s.status == s.STARTED
container = s.service.container
assert container.labels["traefik.enable"] == "true"
assert container.labels[
f"traefik.http.services.{s.hostname}-http.loadbalancer.server.port"] == str(
s.workstation_image.http_port)
assert container.labels[
f"traefik.http.services.{s.hostname}-websocket.loadbalancer.server.port"] == str(
s.workstation_image.websocket_port)
networks = container.attrs.get("NetworkSettings")["Networks"]
assert len(networks) == 1
assert settings.WORKSTATIONS_NETWORK_NAME in networks
with capture_on_commit_callbacks(execute=True):
s.user_finished = True
s.save()
with pytest.raises(NotFound):
# noinspection PyStatementEffect
s.service.container
finally:
stop_all_sessions()
def test_session_redirect(client):
user = UserFactory()
wsi = WorkstationImageFactory(
workstation__title=settings.DEFAULT_WORKSTATION_SLUG, ready=True)
wsi.workstation.add_user(user=user)
response = get_view_for_user(
client=client,
method=client.post,
viewname="workstations:default-session-create",
user=user,
data={"region": "eu-central-1"},
)
assert response.status_code == 302
response = get_view_for_user(client=client, user=user, url=response.url)
assert response.status_code == 200
def test_correct_session_stopped(http_image, docker_client, settings):
path, sha256 = http_image
wsi = WorkstationImageFactory(image__from_path=path,
image_sha256=sha256,
ready=True)
# Execute the celery in place
settings.task_eager_propagates = (True, )
settings.task_always_eager = (True, )
try:
with capture_on_commit_callbacks(execute=True):
s1, s2 = (
SessionFactory(workstation_image=wsi),
SessionFactory(workstation_image=wsi),
)
assert s1.service.container
assert s2.service.container
s2.refresh_from_db()
auth_token_pk = s2.auth_token.pk
with capture_on_commit_callbacks(execute=True):
s2.user_finished = True
s2.save()
assert s1.service.container
with pytest.raises(NotFound):
# noinspection PyStatementEffect
s2.service.container
with pytest.raises(ObjectDoesNotExist):
# auth token should be deleted when the service is stopped
AuthToken.objects.get(pk=auth_token_pk)
finally:
stop_all_sessions()