Merge branch 'ol-test-allocation-endpoints' into ol-dbus-async-v2

olethanh · olethanh · commit 82398b9a0ddb · 2024-04-30T15:15:18.000+02:00
diff --git a/.github/workflows/test-on-droplets-matrix.yml b/.github/workflows/test-on-droplets-matrix.yml
@@ -238,6 +238,14 @@ jobs:
             -d '{"persistent_vms": [], "instances": ["${{ matrix.check_vm.item_hash }}"]}' \
             "http://${DROPLET_IPV4}:4020/control/allocations"
 
+      - name: Get system usage
+        run: |
+          export DROPLET_IPV4="$(doctl compute droplet get aleph-vm-ci-${{ matrix.os_config.alias }}-${{ matrix.check_vm.alias }} --output json | ./.github/scripts/extract_droplet_ipv4.py)"
+          curl -X GET -H "Content-Type: application/json" \
+            -H "X-Auth-Signature: test" \
+            "http://${DROPLET_IPV4}:4020/about/usage/system"
+
+
       - name: Export aleph logs
         if: always()
         run: |
diff --git a/src/aleph/vm/orchestrator/cli.py b/src/aleph/vm/orchestrator/cli.py
@@ -251,8 +251,8 @@ async def run_instances(instances: list[ItemHash]) -> None:
     # Watching for updates on this instance will therefore not work.
     pubsub: Optional[PubSub] = None
 
-    await asyncio.gather(*[start_instance(instance_id, pubsub, pool) for instance_id in instances])
 
+    await asyncio.gather(*[start_instance(instance_id, pubsub, pool) for instance_id in instances])
     await asyncio.Event().wait()  # wait forever
 
 
diff --git a/src/aleph/vm/orchestrator/resources.py b/src/aleph/vm/orchestrator/resources.py
@@ -88,7 +88,7 @@ def get_machine_properties() -> MachineProperties:
     return MachineProperties(
         cpu=CpuProperties(
             architecture=cpu_info.get("raw_arch_string", cpu_info.get("arch_string_raw")),
-            vendor=cpu_info["vendor_id_raw"],
+            vendor=cpu_info.get("vendor_id", cpu_info.get("vendor_id_raw")),
         ),
     )
 
diff --git a/src/aleph/vm/orchestrator/supervisor.py b/src/aleph/vm/orchestrator/supervisor.py
@@ -74,62 +74,63 @@ async def http_not_found(request: web.Request):
     return web.HTTPNotFound()
 
 
-app = web.Application(middlewares=[server_version_middleware])
-cors = setup(
-    app,
-    defaults={
-        "*": ResourceOptions(
-            allow_credentials=True,
-            expose_headers="*",
-            allow_headers="*",
-        )
-    },
-)
+def setup_webapp():
+    app = web.Application(middlewares=[server_version_middleware])
+    cors = setup(
+        app,
+        defaults={
+            "*": ResourceOptions(
+                allow_credentials=True,
+                expose_headers="*",
+                allow_headers="*",
+            )
+        },
+    )
 
-# Routes that need CORS enabled
-cors_routes = [
-    # /about APIs return information about the VM Orchestrator
-    web.get("/about/login", about_login),
-    web.get("/about/executions/list", list_executions),
-    web.get("/about/executions/details", about_executions),
-    web.get("/about/executions/records", about_execution_records),
-    web.get("/about/usage/system", about_system_usage),
-    web.get("/about/config", about_config),
-    # /control APIs are used to control the VMs and access their logs
-    web.post("/control/allocation/notify", notify_allocation),
-    web.get("/control/machine/{ref}/logs", stream_logs),
-    web.post("/control/machine/{ref}/expire", operate_expire),
-    web.post("/control/machine/{ref}/stop", operate_stop),
-    web.post("/control/machine/{ref}/erase", operate_erase),
-    web.post("/control/machine/{ref}/reboot", operate_reboot),
-    # /status APIs are used to check that the VM Orchestrator is running properly
-    web.get("/status/check/fastapi", status_check_fastapi),
-    web.get("/status/check/fastapi/legacy", status_check_fastapi_legacy),
-    web.get("/status/check/host", status_check_host),
-    web.get("/status/check/version", status_check_version),
-    web.get("/status/check/ipv6", status_check_ipv6),
-    web.get("/status/config", status_public_config),
-]
-routes = app.add_routes(cors_routes)
-for route in routes:
-    cors.add(route)
-
-
-# Routes that don't need CORS enabled
-other_routes = [
-    # /control APIs are used to control the VMs and access their logs
-    web.post("/control/allocations", update_allocations),
-    # Raise an HTTP Error 404 if attempting to access an unknown URL within these paths.
-    web.get("/about/{suffix:.*}", http_not_found),
-    web.get("/control/{suffix:.*}", http_not_found),
-    web.get("/status/{suffix:.*}", http_not_found),
-    # /static is used to serve static files
-    web.static("/static", Path(__file__).parent / "views/static"),
-    # /vm is used to launch VMs on-demand
-    web.route("*", "/vm/{ref}{suffix:.*}", run_code_from_path),
-    web.route("*", "/{suffix:.*}", run_code_from_hostname),
-]
-app.add_routes(other_routes)
+    # Routes that need CORS enabled
+    cors_routes = [
+        # /about APIs return information about the VM Orchestrator
+        web.get("/about/login", about_login),
+        web.get("/about/executions/list", list_executions),
+        web.get("/about/executions/details", about_executions),
+        web.get("/about/executions/records", about_execution_records),
+        web.get("/about/usage/system", about_system_usage),
+        web.get("/about/config", about_config),
+        # /control APIs are used to control the VMs and access their logs
+        web.post("/control/allocation/notify", notify_allocation),
+        web.get("/control/machine/{ref}/logs", stream_logs),
+        web.post("/control/machine/{ref}/expire", operate_expire),
+        web.post("/control/machine/{ref}/stop", operate_stop),
+        web.post("/control/machine/{ref}/erase", operate_erase),
+        web.post("/control/machine/{ref}/reboot", operate_reboot),
+        # /status APIs are used to check that the VM Orchestrator is running properly
+        web.get("/status/check/fastapi", status_check_fastapi),
+        web.get("/status/check/fastapi/legacy", status_check_fastapi_legacy),
+        web.get("/status/check/host", status_check_host),
+        web.get("/status/check/version", status_check_version),
+        web.get("/status/check/ipv6", status_check_ipv6),
+        web.get("/status/config", status_public_config),
+    ]
+    routes = app.add_routes(cors_routes)
+    for route in routes:
+        cors.add(route)
+
+    # Routes that don't need CORS enabled
+    other_routes = [
+        # /control APIs are used to control the VMs and access their logs
+        web.post("/control/allocations", update_allocations),
+        # Raise an HTTP Error 404 if attempting to access an unknown URL within these paths.
+        web.get("/about/{suffix:.*}", http_not_found),
+        web.get("/control/{suffix:.*}", http_not_found),
+        web.get("/status/{suffix:.*}", http_not_found),
+        # /static is used to serve static files
+        web.static("/static", Path(__file__).parent / "views/static"),
+        # /vm is used to launch VMs on-demand
+        web.route("*", "/vm/{ref}{suffix:.*}", run_code_from_path),
+        web.route("*", "/{suffix:.*}", run_code_from_hostname),
+    ]
+    app.add_routes(other_routes)
+    return app
 
 
 async def stop_all_vms(app: web.Application):
@@ -153,6 +154,7 @@ def run():
 
     # Require a random token to access /about APIs
     secret_token = token_urlsafe(nbytes=32)
+    app = setup_webapp()
     # Store app singletons. Note that app["pubsub"] will also be created.
     app["secret_token"] = secret_token
     app["vm_pool"] = pool
diff --git a/src/aleph/vm/orchestrator/views/__init__.py b/src/aleph/vm/orchestrator/views/__init__.py
@@ -341,6 +341,12 @@ def authenticate_api_request(request: web.Request) -> bool:
 
 
 async def update_allocations(request: web.Request):
+    """Main entry for the start of persistence VM and instance, called by the CCN,
+
+
+    auth via the SETTINGS.ALLOCATION_TOKEN_HASH  sent in header X-Auth-Signature.
+    Receive a list of vm and instance that should be present and then match that state by stopping and launching VMs
+    """
     if not authenticate_api_request(request):
         return web.HTTPUnauthorized(text="Authentication token received is invalid")
 
diff --git a/tests/supervisor/test_execution.py b/tests/supervisor/test_execution.py
@@ -58,11 +58,13 @@ async def test_create_execution():
 
 
 @pytest.mark.asyncio
-async def test_create_execution_online():
+async def test_create_execution_online(vm_hash: ItemHash = None):
     """
     Create a new VM execution without building it locally and check that it starts properly.
     """
 
+    vm_hash = vm_hash or settings.CHECK_FASTAPI_VM_ID
+
     # Ensure that the settings are correct and required files present.
     settings.setup()
     settings.check()
@@ -71,7 +73,6 @@ async def test_create_execution_online():
     engine = metrics.setup_engine()
     await metrics.create_tables(engine)
 
-    vm_hash = ItemHash("3fc0aa9569da840c43e7bd2033c3c580abb46b007527d6d20f2d4e98e867f7af")
     message = await get_message(ref=vm_hash)
 
     execution = VmExecution(
@@ -93,3 +94,11 @@ async def test_create_execution_online():
 
     await execution.start()
     await execution.stop()
+
+
+@pytest.mark.asyncio
+async def test_create_execution_legacy():
+    """
+    Create a new VM execution based on the legacy FastAPI check and ensure that it starts properly.
+    """
+    await test_create_execution_online(vm_hash=settings.LEGACY_CHECK_FASTAPI_VM_ID)
diff --git a/tests/supervisor/test_views.py b/tests/supervisor/test_views.py
@@ -2,12 +2,13 @@
 from aiohttp import web
 
 from aleph.vm.conf import settings
-from aleph.vm.orchestrator.supervisor import app
+from aleph.vm.orchestrator.supervisor import setup_webapp
 
 
 @pytest.mark.asyncio
 async def test_allocation_fails_on_invalid_item_hash(aiohttp_client):
     """Test that the allocation endpoint fails when an invalid item_hash is provided."""
+    app = setup_webapp()
     client = await aiohttp_client(app)
     settings.ALLOCATION_TOKEN_HASH = "9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08"  # = "test"
     response: web.Response = await client.post(
@@ -28,12 +29,65 @@ async def test_allocation_fails_on_invalid_item_hash(aiohttp_client):
 
 @pytest.mark.asyncio
 async def test_system_usage(aiohttp_client):
-    """Test that the allocation endpoint fails when an invalid item_hash is provided."""
+    """Test that the usage system endpoints responds. No auth needed"""
+    app = setup_webapp()
     client = await aiohttp_client(app)
-    settings.ALLOCATION_TOKEN_HASH = "9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08"  # = "test"
     response: web.Response = await client.get("/about/usage/system")
     assert response.status == 200
     # check if it is valid json
     resp = await response.json()
     assert "cpu" in resp
     assert resp["cpu"]["count"] > 0
+
+
+@pytest.mark.asyncio
+async def test_allocation_invalid_auth_token(aiohttp_client):
+    """Test that the allocation endpoint fails when an invalid auth token is provided."""
+    settings.ALLOCATION_TOKEN_HASH = "9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08"  # = "test"
+    app = setup_webapp()
+    client = await aiohttp_client(app)
+    response = await client.post(
+        "/control/allocations",
+        json={"persistent_vms": []},
+        headers={"X-Auth-Signature": "notTest"},
+    )
+    assert response.status == 401
+    assert await response.text() == "Authentication token received is invalid"
+
+
+@pytest.mark.asyncio
+async def test_allocation_missing_auth_token(aiohttp_client):
+    """Test that the allocation endpoint fails when auth token is not provided."""
+    app = setup_webapp()
+    client = await aiohttp_client(app)
+    response: web.Response = await client.post(
+        "/control/allocations",
+        json={"persistent_vms": []},
+    )
+    assert response.status == 401
+    assert await response.text() == "Authentication token is missing"
+
+
+@pytest.mark.asyncio
+async def test_allocation_valid_token(aiohttp_client):
+    """Test that the allocation endpoint fails when an invalid auth is provided.
+
+    This is a very simple test that don't start or stop any VM so the mock is minimal"""
+
+    class FakeVmPool:
+        def get_persistent_executions(self):
+            return []
+
+    settings.ALLOCATION_TOKEN_HASH = "9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08"  # = "test"
+    app = setup_webapp()
+    app["vm_pool"] = FakeVmPool()
+    app["pubsub"] = FakeVmPool()
+    client = await aiohttp_client(app)
+
+    response: web.Response = await client.post(
+        "/control/allocations",
+        json={"persistent_vms": []},
+        headers={"X-Auth-Signature": "test"},
+    )
+    assert response.status == 200
+    assert await response.json() == {"success": True, "successful": [], "failing": [], "errors": {}}

Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,7 @@ def get_machine_properties() -> MachineProperties:`
`88`	`88`	`return MachineProperties(`
`89`	`89`	`cpu=CpuProperties(`
`90`	`90`	`architecture=cpu_info.get("raw_arch_string", cpu_info.get("arch_string_raw")),`
`91`		`- vendor=cpu_info["vendor_id_raw"],`
	`91`	`+ vendor=cpu_info.get("vendor_id", cpu_info.get("vendor_id_raw")),`
`92`	`92`	`),`
`93`	`93`	`)`
`94`	`94`