258 lines
7.1 KiB
Python
258 lines
7.1 KiB
Python
"""Main CLI module for TestDataGen."""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import click
|
|
|
|
from testdatagen.formatters.csv_formatter import CSVFormatter
|
|
from testdatagen.formatters.json_formatter import JSONFormatter
|
|
from testdatagen.formatters.sql_formatter import SQLFormatter
|
|
from testdatagen.generators.json_schema_generator import JSONSchemaGenerator
|
|
|
|
|
|
@click.group()
|
|
@click.version_option(version="0.1.0")
|
|
def main():
|
|
"""TestDataGen - Generate realistic test data from schemas and types."""
|
|
pass
|
|
|
|
|
|
@main.command()
|
|
@click.option(
|
|
"--schema", "-s",
|
|
type=click.Path(exists=True, file_okay=True, dir_okay=False),
|
|
required=True,
|
|
help="Path to JSON Schema file"
|
|
)
|
|
@click.option(
|
|
"--count", "-n",
|
|
type=int,
|
|
default=10,
|
|
help="Number of records to generate (default: 10)"
|
|
)
|
|
@click.option(
|
|
"--format", "-f",
|
|
type=click.Choice(["json", "csv", "sql"], case_sensitive=False),
|
|
default="json",
|
|
help="Output format (default: json)"
|
|
)
|
|
@click.option(
|
|
"--seed",
|
|
type=int,
|
|
default=None,
|
|
help="Random seed for reproducibility"
|
|
)
|
|
@click.option(
|
|
"--table",
|
|
type=str,
|
|
default="generated_table",
|
|
help="Table name for SQL output (default: generated_table)"
|
|
)
|
|
@click.option(
|
|
"--indent",
|
|
type=int,
|
|
default=None,
|
|
help="Indentation level for JSON output (default: None)"
|
|
)
|
|
def generate(schema, count, format, seed, table, indent):
|
|
"""Generate test data from a JSON Schema file."""
|
|
try:
|
|
schema_path = Path(schema)
|
|
with open(schema_path, "r") as f:
|
|
import json
|
|
schema_data = json.load(f)
|
|
|
|
generator = JSONSchemaGenerator(seed=seed)
|
|
records = generator.generate(schema_data, count=count)
|
|
|
|
if format.lower() == "json":
|
|
formatter = JSONFormatter(indent=indent)
|
|
elif format.lower() == "csv":
|
|
formatter = CSVFormatter()
|
|
elif format.lower() == "sql":
|
|
formatter = SQLFormatter(table_name=table)
|
|
else:
|
|
click.echo(f"Error: Unsupported format '{format}'", err=True)
|
|
sys.exit(1)
|
|
|
|
output = formatter.format(records)
|
|
click.echo(output)
|
|
|
|
except json.JSONDecodeError as e:
|
|
click.echo(f"Error: Invalid JSON in schema file: {e}", err=True)
|
|
sys.exit(1)
|
|
except FileNotFoundError:
|
|
click.echo(f"Error: Schema file not found: {schema}", err=True)
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
click.echo(f"Error: {e}", err=True)
|
|
sys.exit(1)
|
|
|
|
|
|
@main.command()
|
|
@click.option(
|
|
"--input", "-i",
|
|
type=click.Path(exists=True, file_okay=True, dir_okay=False),
|
|
required=True,
|
|
help="Path to TypeScript file"
|
|
)
|
|
@click.option(
|
|
"--count", "-n",
|
|
type=int,
|
|
default=10,
|
|
help="Number of records to generate (default: 10)"
|
|
)
|
|
@click.option(
|
|
"--format", "-f",
|
|
type=click.Choice(["json", "csv", "sql"], case_sensitive=False),
|
|
default="json",
|
|
help="Output format (default: json)"
|
|
)
|
|
@click.option(
|
|
"--seed",
|
|
type=int,
|
|
default=None,
|
|
help="Random seed for reproducibility"
|
|
)
|
|
@click.option(
|
|
"--table",
|
|
type=str,
|
|
default="generated_table",
|
|
help="Table name for SQL output (default: generated_table)"
|
|
)
|
|
def from_ts(input, count, format, seed, table):
|
|
"""Generate test data from a TypeScript type definition."""
|
|
try:
|
|
import subprocess
|
|
result = subprocess.run(
|
|
["npx", "tsc", "--declaration", "--emitDeclarationOnly", "--jsonSchemaManifest", input],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=30
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
click.echo(f"Error: TypeScript compilation failed: {result.stderr}", err=True)
|
|
sys.exit(1)
|
|
|
|
schema_path = Path(input).with_suffix(".json")
|
|
if not schema_path.exists():
|
|
click.echo("Error: Could not generate schema from TypeScript file", err=True)
|
|
sys.exit(1)
|
|
|
|
with open(schema_path, "r") as f:
|
|
import json
|
|
schema_data = json.load(f)
|
|
|
|
generator = JSONSchemaGenerator(seed=seed)
|
|
records = generator.generate(schema_data, count=count)
|
|
|
|
if format.lower() == "json":
|
|
formatter = JSONFormatter()
|
|
elif format.lower() == "csv":
|
|
formatter = CSVFormatter()
|
|
elif format.lower() == "sql":
|
|
formatter = SQLFormatter(table_name=table)
|
|
else:
|
|
click.echo(f"Error: Unsupported format '{format}'", err=True)
|
|
sys.exit(1)
|
|
|
|
output = formatter.format(records)
|
|
click.echo(output)
|
|
|
|
except FileNotFoundError:
|
|
click.echo("Error: TypeScript file not found", err=True)
|
|
sys.exit(1)
|
|
except subprocess.TimeoutExpired:
|
|
click.echo("Error: TypeScript compilation timed out", err=True)
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
click.echo(f"Error: {e}", err=True)
|
|
sys.exit(1)
|
|
|
|
|
|
@main.command()
|
|
@click.option(
|
|
"--input", "-i",
|
|
type=click.Path(exists=True, file_okay=True, dir_okay=False),
|
|
required=True,
|
|
help="Path to sample data file (JSON or CSV)"
|
|
)
|
|
@click.option(
|
|
"--count", "-n",
|
|
type=int,
|
|
default=10,
|
|
help="Number of records to generate (default: 10)"
|
|
)
|
|
@click.option(
|
|
"--format", "-f",
|
|
type=click.Choice(["json", "csv", "sql"], case_sensitive=False),
|
|
default="json",
|
|
help="Output format (default: json)"
|
|
)
|
|
@click.option(
|
|
"--seed",
|
|
type=int,
|
|
default=None,
|
|
help="Random seed for reproducibility"
|
|
)
|
|
@click.option(
|
|
"--table",
|
|
type=str,
|
|
default="generated_table",
|
|
help="Table name for SQL output (default: generated_table)"
|
|
)
|
|
def from_sample(input, count, format, seed, table):
|
|
"""Generate test data from a sample data file."""
|
|
try:
|
|
input_path = Path(input)
|
|
|
|
with open(input_path, "r") as f:
|
|
import json
|
|
sample_data = json.load(f)
|
|
|
|
try:
|
|
from genson import SchemaBuilder
|
|
except ImportError:
|
|
click.echo("Error: genson not installed. Run: pip install genson", err=True)
|
|
sys.exit(1)
|
|
|
|
builder = SchemaBuilder()
|
|
if isinstance(sample_data, list):
|
|
for item in sample_data:
|
|
builder.add_object(item)
|
|
else:
|
|
builder.add_object(sample_data)
|
|
|
|
schema_data = builder.to_schema()
|
|
|
|
generator = JSONSchemaGenerator(seed=seed)
|
|
records = generator.generate(schema_data, count=count)
|
|
|
|
if format.lower() == "json":
|
|
formatter = JSONFormatter()
|
|
elif format.lower() == "csv":
|
|
formatter = CSVFormatter()
|
|
elif format.lower() == "sql":
|
|
formatter = SQLFormatter(table_name=table)
|
|
else:
|
|
click.echo(f"Error: Unsupported format '{format}'", err=True)
|
|
sys.exit(1)
|
|
|
|
output = formatter.format(records)
|
|
click.echo(output)
|
|
|
|
except json.JSONDecodeError:
|
|
click.echo(f"Error: Invalid JSON in sample file: {input}", err=True)
|
|
sys.exit(1)
|
|
except FileNotFoundError:
|
|
click.echo(f"Error: Sample file not found: {input}", err=True)
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
click.echo(f"Error: {e}", err=True)
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |