fix(Scan): PRWLR-5145 Update resource storing when performing a scan (#67)

* chore(Resource): PRWLR-5145 remove region from unique constraint

* feat(Resource): PRWLR-5145 update how resources are stored during a scan

* test(Task): PRWLR-5145 add unit tests for scan task resources
This commit is contained in:
Víctor Fernández Poyatos
2024-11-04 11:05:02 +01:00
committed by GitHub
parent 31c035eb52
commit 3de2c47c56
4 changed files with 158 additions and 7 deletions
+1 -1
View File
@@ -864,7 +864,7 @@ class Migration(migrations.Migration):
migrations.AddConstraint(
model_name="resource",
constraint=models.UniqueConstraint(
fields=("tenant_id", "provider_id", "uid", "region"),
fields=("tenant_id", "provider_id", "uid"),
name="unique_resources_by_provider",
),
),
+1 -1
View File
@@ -432,7 +432,7 @@ class Resource(RowLevelSecurityProtectedModel):
constraints = [
models.UniqueConstraint(
fields=("tenant_id", "provider_id", "uid", "region"),
fields=("tenant_id", "provider_id", "uid"),
name="unique_resources_by_provider",
),
RowLevelSecurityConstraint(
+12 -4
View File
@@ -108,14 +108,22 @@ def _store_resources(
"""
with tenant_transaction(tenant_id):
resource_instance, _ = Resource.objects.get_or_create(
resource_instance, created = Resource.objects.get_or_create(
tenant_id=tenant_id,
provider=provider_instance,
uid=finding.resource_uid,
region=finding.region,
service=finding.service_name,
type=finding.resource_type,
defaults={
"region": finding.region,
"service": finding.service_name,
"type": finding.resource_type,
},
)
if not created:
resource_instance.region = finding.region
resource_instance.service = finding.service_name
resource_instance.type = finding.resource_type
resource_instance.save()
with tenant_transaction(tenant_id):
tags = [
ResourceTag.objects.get_or_create(
+144 -1
View File
@@ -3,7 +3,11 @@ from unittest.mock import patch, MagicMock
import pytest
from api.models import StateChoices, StatusChoices, Severity, Finding, Resource
from tasks.jobs.scan import perform_prowler_scan, _create_finding_delta
from tasks.jobs.scan import (
perform_prowler_scan,
_create_finding_delta,
_store_resources,
)
@pytest.mark.django_db
@@ -112,3 +116,142 @@ class TestPerformScan:
)
def test_create_finding_delta(self, last_status, new_status, expected_delta):
assert _create_finding_delta(last_status, new_status) == expected_delta
@patch("tasks.jobs.scan.ResourceTag.objects.get_or_create")
@patch("tasks.jobs.scan.Resource.objects.get_or_create")
@patch("api.db_utils.tenant_transaction")
def test_store_resources_new_resource(
self,
mock_tenant_transaction,
mock_get_or_create_resource,
mock_get_or_create_tag,
):
tenant_id = "tenant123"
provider_instance = MagicMock()
provider_instance.id = "provider456"
finding = MagicMock()
finding.resource_uid = "resource_uid_123"
finding.region = "us-west-1"
finding.service_name = "service_name"
finding.resource_type = "resource_type"
finding.resource_tags = {"tag1": "value1", "tag2": "value2"}
resource_instance = MagicMock()
mock_get_or_create_resource.return_value = (resource_instance, True)
tag_instance = MagicMock()
mock_get_or_create_tag.return_value = (tag_instance, True)
resource, resource_uid_tuple = _store_resources(
finding, tenant_id, provider_instance
)
mock_get_or_create_resource.assert_called_once_with(
tenant_id=tenant_id,
provider=provider_instance,
uid=finding.resource_uid,
defaults={
"region": finding.region,
"service": finding.service_name,
"type": finding.resource_type,
},
)
assert resource == resource_instance
assert resource_uid_tuple == (resource_instance.uid, resource_instance.region)
@patch("tasks.jobs.scan.ResourceTag.objects.get_or_create")
@patch("tasks.jobs.scan.Resource.objects.get_or_create")
@patch("api.db_utils.tenant_transaction")
def test_store_resources_existing_resource(
self,
mock_tenant_transaction,
mock_get_or_create_resource,
mock_get_or_create_tag,
):
tenant_id = "tenant123"
provider_instance = MagicMock()
provider_instance.id = "provider456"
finding = MagicMock()
finding.resource_uid = "resource_uid_123"
finding.region = "us-west-2"
finding.service_name = "new_service"
finding.resource_type = "new_type"
finding.resource_tags = {"tag1": "value1", "tag2": "value2"}
resource_instance = MagicMock()
resource_instance.region = "us-west-1"
resource_instance.service = "old_service"
resource_instance.type = "old_type"
mock_get_or_create_resource.return_value = (resource_instance, False)
tag_instance = MagicMock()
mock_get_or_create_tag.return_value = (tag_instance, True)
resource, resource_uid_tuple = _store_resources(
finding, tenant_id, provider_instance
)
mock_get_or_create_resource.assert_called_once_with(
tenant_id=tenant_id,
provider=provider_instance,
uid=finding.resource_uid,
defaults={
"region": finding.region,
"service": finding.service_name,
"type": finding.resource_type,
},
)
assert resource_instance.region == finding.region
assert resource_instance.service == finding.service_name
assert resource_instance.type == finding.resource_type
assert resource == resource_instance
assert resource_uid_tuple == (resource_instance.uid, resource_instance.region)
@patch("tasks.jobs.scan.ResourceTag.objects.get_or_create")
@patch("tasks.jobs.scan.Resource.objects.get_or_create")
@patch("api.db_utils.tenant_transaction")
def test_store_resources_with_tags(
self,
mock_tenant_transaction,
mock_get_or_create_resource,
mock_get_or_create_tag,
):
tenant_id = "tenant123"
provider_instance = MagicMock()
provider_instance.id = "provider456"
finding = MagicMock()
finding.resource_uid = "resource_uid_123"
finding.region = "us-west-1"
finding.service_name = "service_name"
finding.resource_type = "resource_type"
finding.resource_tags = {"tag1": "value1", "tag2": "value2"}
resource_instance = MagicMock()
mock_get_or_create_resource.return_value = (resource_instance, True)
tag_instance_1 = MagicMock()
tag_instance_2 = MagicMock()
mock_get_or_create_tag.side_effect = [
(tag_instance_1, True),
(tag_instance_2, True),
]
resource, resource_uid_tuple = _store_resources(
finding, tenant_id, provider_instance
)
mock_get_or_create_tag.assert_any_call(
tenant_id=tenant_id, key="tag1", value="value1"
)
mock_get_or_create_tag.assert_any_call(
tenant_id=tenant_id, key="tag2", value="value2"
)
resource_instance.upsert_or_delete_tags.assert_called_once_with(
tags=[tag_instance_1, tag_instance_2]
)
assert resource == resource_instance
assert resource_uid_tuple == (resource_instance.uid, resource_instance.region)