#!/usr/bin/env python # # Copyright (C) 2021 GNS3 Technologies Inc. # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU General Public License as published by # the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with this program. If not, see . import os import pytest import hashlib from sqlalchemy.ext.asyncio import AsyncSession from fastapi import FastAPI, status from httpx import AsyncClient from gns3server.db.repositories.images import ImagesRepository pytestmark = pytest.mark.asyncio @pytest.fixture def iou_32_bit_image(tmpdir) -> str: """ Create a fake IOU image on disk """ path = os.path.join(tmpdir, "iou_32bit.bin") with open(path, "wb+") as f: f.write(b'\x7fELF\x01\x01\x01') return path @pytest.fixture def iou_64_bit_image(tmpdir) -> str: """ Create a fake IOU image on disk """ path = os.path.join(tmpdir, "iou_64bit.bin") with open(path, "wb+") as f: f.write(b'\x7fELF\x02\x01\x01') return path @pytest.fixture def ios_image(tmpdir) -> str: """ Create a fake IOS image on disk """ path = os.path.join(tmpdir, "ios.bin") with open(path, "wb+") as f: f.write(b'\x7fELF\x01\x02\x01') return path @pytest.fixture def qcow2_image(tmpdir) -> str: """ Create a fake Qemu qcow2 image on disk """ path = os.path.join(tmpdir, "image.qcow2") with open(path, "wb+") as f: f.write(b'QFI\xfb') return path @pytest.fixture def invalid_image(tmpdir) -> str: """ Create a fake invalid image on disk """ path = os.path.join(tmpdir, "invalid_image.bin") with open(path, "wb+") as f: f.write(b'\x01\x01\x01\x01') return path @pytest.fixture def empty_image(tmpdir) -> str: """ Create a fake empty image on disk """ path = os.path.join(tmpdir, "empty_image.bin") with open(path, "wb+") as f: f.write(b'') return path class TestImageRoutes: @pytest.mark.parametrize( "image_type, fixture_name, valid_request", ( ("iou", "iou_32_bit_image", True), ("iou", "iou_64_bit_image", True), ("iou", "invalid_image", False), ("ios", "ios_image", True), ("ios", "invalid_image", False), ("qemu", "qcow2_image", True), ("qemu", "empty_image", False), ("wrong_type", "qcow2_image", False), ), ) async def test_upload_image( self, app: FastAPI, client: AsyncClient, images_dir: str, image_type: str, fixture_name: str, valid_request: bool, request ) -> None: image_path = request.getfixturevalue(fixture_name) image_name = os.path.basename(image_path) image_checksum = hashlib.md5() with open(image_path, "rb") as f: image_data = f.read() image_checksum.update(image_data) response = await client.post( app.url_path_for("upload_image", image_name=image_name), params={"image_type": image_type}, content=image_data) if valid_request: assert response.status_code == status.HTTP_201_CREATED assert response.json()["filename"] == image_name assert response.json()["checksum"] == image_checksum.hexdigest() assert os.path.exists(os.path.join(images_dir, image_type.upper(), image_name)) else: assert response.status_code != status.HTTP_201_CREATED async def test_image_list(self, app: FastAPI, client: AsyncClient) -> None: response = await client.get(app.url_path_for("get_images")) assert response.status_code == status.HTTP_200_OK assert len(response.json()) == 4 # 4 valid images uploaded before async def test_image_get(self, app: FastAPI, client: AsyncClient, qcow2_image: str) -> None: image_name = os.path.basename(qcow2_image) response = await client.get(app.url_path_for("get_image", image_name=image_name)) assert response.status_code == status.HTTP_200_OK assert response.json()["filename"] == image_name async def test_same_image_cannot_be_uploaded(self, app: FastAPI, client: AsyncClient, qcow2_image: str) -> None: image_name = os.path.basename(qcow2_image) with open(qcow2_image, "rb") as f: image_data = f.read() response = await client.post( app.url_path_for("upload_image", image_name=image_name), params={"image_type": "qemu"}, content=image_data) assert response.status_code == status.HTTP_400_BAD_REQUEST async def test_image_delete(self, app: FastAPI, client: AsyncClient, images_dir: str, qcow2_image: str) -> None: image_name = os.path.basename(qcow2_image) response = await client.delete(app.url_path_for("delete_image", image_name=image_name)) assert response.status_code == status.HTTP_204_NO_CONTENT async def test_not_found_image(self, app: FastAPI, client: AsyncClient, qcow2_image: str) -> None: image_name = os.path.basename(qcow2_image) response = await client.get(app.url_path_for("get_image", image_name=image_name)) assert response.status_code == status.HTTP_404_NOT_FOUND async def test_image_deleted_on_disk(self, app: FastAPI, client: AsyncClient, images_dir: str, qcow2_image: str) -> None: image_name = os.path.basename(qcow2_image) with open(qcow2_image, "rb") as f: image_data = f.read() response = await client.post( app.url_path_for("upload_image", image_name=image_name), params={"image_type": "qemu"}, content=image_data) assert response.status_code == status.HTTP_201_CREATED response = await client.delete(app.url_path_for("delete_image", image_name=image_name)) assert response.status_code == status.HTTP_204_NO_CONTENT assert not os.path.exists(os.path.join(images_dir, "QEMU", image_name)) async def test_prune_images(self, app: FastAPI, client: AsyncClient, db_session: AsyncSession) -> None: response = await client.post(app.url_path_for("prune_images")) assert response.status_code == status.HTTP_204_NO_CONTENT images_repo = ImagesRepository(db_session) images_in_db = await images_repo.get_images() assert len(images_in_db) == 0