Skip to content

Commit

Permalink
Stop scan crashing when casting none to int
Browse files Browse the repository at this point in the history
  • Loading branch information
ostorlab committed Dec 17, 2024
1 parent 3328dc7 commit 35e923e
Show file tree
Hide file tree
Showing 2 changed files with 72 additions and 2 deletions.
8 changes: 6 additions & 2 deletions src/ostorlab/runtimes/local/runtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -309,9 +309,13 @@ def stop(self, scan_id: str) -> None:
networks = self._docker_client.networks.list()
for network in networks:
network_labels = network.attrs["Labels"]
if network_labels is None:
continue
universe = network_labels.get("ostorlab.universe")
if (
network_labels is not None
and int(network_labels.get("ostorlab.universe")) == scan_id
universe is not None
and network_labels is not None
and int(universe) == scan_id
):
logger.info("removing network %s", network_labels)
stopped_network.append(network)
Expand Down
66 changes: 66 additions & 0 deletions tests/runtimes/local/runtime_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import docker
import pytest
from docker.models import services as services_model
from docker.models import networks as networks_model
from pytest_mock import plugin

import ostorlab
Expand Down Expand Up @@ -364,3 +365,68 @@ def testCheckServicesMethod_whenServicesAreStopped_shouldExit(
assert exit_mock.call_count == 1
exit_with = exit_mock.call_args_list[0][0][0]
assert exit_with == 0


@pytest.mark.docker
def testRuntimeScanStop_whenUnrelatedNetworks_RemovesScanServiceWithoutCrash(
mocker, db_engine_path
):
"""Unittest for the scan stop method when there are networks not related to the scan, the process shouldn't crash"""
mocker.patch.object(models, "ENGINE_URL", db_engine_path)
create_scan_db = models.Scan.create("test")

Check warning on line 376 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L375-L376

Added lines #L375 - L376 were not covered by tests

def docker_services():

Check warning on line 378 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L378

Added line #L378 was not covered by tests
"""Method for mocking the services list response."""
with models.Database() as session:
scan = session.query(models.Scan).first()
services = [

Check warning on line 382 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L380-L382

Added lines #L380 - L382 were not covered by tests
{
"ID": "0099i5n1y3gycuekvksyqyxav",
"CreatedAt": "2021-12-27T13:37:02.795789947Z",
"Spec": {"Labels": {"ostorlab.universe": scan.id}},
},
{
"ID": "0099i5n1y3gycuekvksyqyxav",
"CreatedAt": "2021-12-27T13:37:02.795789947Z",
"Spec": {"Labels": {"ostorlab.universe": 9999}},
},
]

return [services_model.Service(attrs=service) for service in services]

Check warning on line 395 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L395

Added line #L395 was not covered by tests

def docker_networks():

Check warning on line 397 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L397

Added line #L397 was not covered by tests
"""Method for mocking the services list response."""
with models.Database() as session:
scan = session.query(models.Scan).first()
networks = [

Check warning on line 401 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L399-L401

Added lines #L399 - L401 were not covered by tests
{
"ID": "0099i5n1y3gycuekvksyqyxav",
"CreatedAt": "2021-12-27T13:37:02.795789947Z",
"Labels": {"ostorlab.universe": scan.id},
},
{
"ID": "0099i5n1y3gycuekvksyqyxav",
"CreatedAt": "2021-12-27T13:37:02.795789947Z",
"Labels": {},
},
]

return [networks_model.Network(attrs=network) for network in networks]

Check warning on line 414 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L414

Added line #L414 was not covered by tests

mocker.patch(

Check warning on line 416 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L416

Added line #L416 was not covered by tests
"docker.DockerClient.services", return_value=services_model.ServiceCollection()
)
mocker.patch("docker.DockerClient.services.list", side_effect=docker_services)
mocker.patch(

Check warning on line 420 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L419-L420

Added lines #L419 - L420 were not covered by tests
"docker.models.networks.NetworkCollection.list", return_value=docker_networks()
)
mocker.patch("docker.models.configs.ConfigCollection.list", return_value=[])
docker_network_remove = mocker.patch("docker.models.networks.Network.remove")
docker_service_remove = mocker.patch(

Check warning on line 425 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L423-L425

Added lines #L423 - L425 were not covered by tests
"docker.models.services.Service.remove", return_value=None
)

local_runtime.LocalRuntime().stop(scan_id=create_scan_db.id)

Check warning on line 429 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L429

Added line #L429 was not covered by tests

docker_service_remove.assert_called_once()
docker_network_remove.assert_called_once()

Check warning on line 432 in tests/runtimes/local/runtime_test.py

View check run for this annotation

Codecov / codecov/patch

tests/runtimes/local/runtime_test.py#L431-L432

Added lines #L431 - L432 were not covered by tests

0 comments on commit 35e923e

Please sign in to comment.