Files
xingrin/tools/seed-api/test_data_generator.py
yyhuni 069527a7f1 feat(backend): implement vulnerability and screenshot snapshot APIs with directories tab reorganization
- Add vulnerability snapshot DTO, handler, repository, and service layer with comprehensive test coverage
- Add screenshot snapshot DTO, handler, repository, and service layer for snapshot management
- Reorganize directories tab from secondary assets navigation to primary navigation in scan history and target layouts
- Update frontend navigation to include FolderSearch icon for directories tab with badge count display
- Add i18n translations for directories tab in English and Chinese messages
- Implement seed data generation tools with Python API client for testing and data population
- Add data generator, error handler, and progress tracking utilities for seed API
- Update target validator to support new snapshot-related validations
- Refactor organization and vulnerability handlers to support snapshot operations
- Add integration tests and property-based tests for vulnerability snapshot functionality
- Update Go module dependencies to support new snapshot features
2026-01-15 10:25:34 +08:00

169 lines
6.2 KiB
Python

"""
Unit tests for Data Generator module.
"""
import pytest
from data_generator import DataGenerator
class TestDataGenerator:
"""Test DataGenerator class."""
def test_generate_organization(self):
"""Test organization generation."""
org = DataGenerator.generate_organization(0)
assert "name" in org
assert "description" in org
assert isinstance(org["name"], str)
assert isinstance(org["description"], str)
assert len(org["name"]) > 0
assert len(org["description"]) > 0
def test_generate_targets(self):
"""Test target generation with correct ratios."""
targets = DataGenerator.generate_targets(100)
# Count types
domain_count = sum(1 for t in targets if t["type"] == "domain")
ip_count = sum(1 for t in targets if t["type"] == "ip")
cidr_count = sum(1 for t in targets if t["type"] == "cidr")
# Check ratios (allow 5% tolerance)
assert 65 <= domain_count <= 75 # 70% ± 5%
assert 15 <= ip_count <= 25 # 20% ± 5%
assert 5 <= cidr_count <= 15 # 10% ± 5%
# Check all have required fields
for target in targets:
assert "name" in target
assert "type" in target
assert target["type"] in ["domain", "ip", "cidr"]
def test_generate_domain_target(self):
"""Test domain target format."""
targets = DataGenerator.generate_targets(10, {"domain": 1.0, "ip": 0, "cidr": 0})
for target in targets:
assert target["type"] == "domain"
assert "." in target["name"]
assert not target["name"].startswith(".")
assert not target["name"].endswith(".")
def test_generate_ip_target(self):
"""Test IP target format."""
targets = DataGenerator.generate_targets(10, {"domain": 0, "ip": 1.0, "cidr": 0})
for target in targets:
assert target["type"] == "ip"
parts = target["name"].split(".")
assert len(parts) == 4
for part in parts:
assert 0 <= int(part) <= 255
def test_generate_cidr_target(self):
"""Test CIDR target format."""
targets = DataGenerator.generate_targets(10, {"domain": 0, "ip": 0, "cidr": 1.0})
for target in targets:
assert target["type"] == "cidr"
assert "/" in target["name"]
ip, mask = target["name"].split("/")
assert int(mask) in [8, 16, 24]
def test_generate_websites(self):
"""Test website generation."""
target = {"name": "example.com", "type": "domain", "id": 1}
websites = DataGenerator.generate_websites(target, 5)
assert len(websites) == 5
for website in websites:
assert "url" in website
assert "title" in website
assert "statusCode" in website
assert "tech" in website
assert isinstance(website["tech"], list)
assert website["url"].startswith("http")
assert "example.com" in website["url"]
def test_generate_subdomains_for_domain(self):
"""Test subdomain generation for domain target."""
target = {"name": "example.com", "type": "domain", "id": 1}
subdomains = DataGenerator.generate_subdomains(target, 5)
assert len(subdomains) == 5
for subdomain in subdomains:
assert "name" in subdomain
assert subdomain["name"].endswith("example.com")
assert subdomain["name"] != "example.com" # Should have prefix
def test_generate_subdomains_for_non_domain(self):
"""Test subdomain generation for non-domain target."""
target = {"name": "192.168.1.1", "type": "ip", "id": 1}
subdomains = DataGenerator.generate_subdomains(target, 5)
assert len(subdomains) == 0 # Should return empty for non-domain
def test_generate_endpoints(self):
"""Test endpoint generation."""
target = {"name": "example.com", "type": "domain", "id": 1}
endpoints = DataGenerator.generate_endpoints(target, 5)
assert len(endpoints) == 5
for endpoint in endpoints:
assert "url" in endpoint
assert "statusCode" in endpoint
assert "tech" in endpoint
assert "matchedGfPatterns" in endpoint
assert isinstance(endpoint["tech"], list)
assert isinstance(endpoint["matchedGfPatterns"], list)
def test_generate_directories(self):
"""Test directory generation."""
target = {"name": "example.com", "type": "domain", "id": 1}
directories = DataGenerator.generate_directories(target, 5)
assert len(directories) == 5
for directory in directories:
assert "url" in directory
assert "status" in directory
assert "contentLength" in directory
assert directory["url"].endswith("/") # Directories end with /
def test_generate_host_ports(self):
"""Test host port generation."""
target = {"name": "example.com", "type": "domain", "id": 1}
host_ports = DataGenerator.generate_host_ports(target, 5)
assert len(host_ports) == 5
for hp in host_ports:
assert "host" in hp
assert "ip" in hp
assert "port" in hp
assert isinstance(hp["port"], int)
assert 1 <= hp["port"] <= 65535
def test_generate_vulnerabilities(self):
"""Test vulnerability generation."""
target = {"name": "example.com", "type": "domain", "id": 1}
vulns = DataGenerator.generate_vulnerabilities(target, 5)
assert len(vulns) == 5
for vuln in vulns:
assert "url" in vuln
assert "vulnType" in vuln
assert "severity" in vuln
assert "cvssScore" in vuln
assert vuln["severity"] in ["critical", "high", "medium", "low", "info"]
assert 0 <= vuln["cvssScore"] <= 10
if __name__ == "__main__":
pytest.main([__file__, "-v"])