test-legoclasses.py

# import functions.lego_classes as lego_classes
from functions.lego_classes import *
import json
import pprint
# Test manually creating some item and components

battery = LegoComponent("nice battery", ComponentCategory.BATTERY, "bat42", 1, 2, 3)
motor = LegoComponent("motor goes brrr", ComponentCategory.MOTOR, "motor", 1, 2, 3)
wheel = LegoComponent("much round wheel", ComponentCategory.WHEEL, "round", 1, 2, 3)

car = LegoAssembly("Car", AggregationLayer.SYSTEM)

chassis = LegoAssembly("Chassis", AggregationLayer.ASSEMBLY)
door1 = LegoAssembly("Door 1", AggregationLayer.SUBASSEMBLY)
door2 = LegoAssembly("Door 2", AggregationLayer.SUBASSEMBLY)
chassis.add_assembly(door1)
chassis.add_assembly(door2)

engine = LegoAssembly("Engine", AggregationLayer.ASSEMBLY)
engine.add_component(motor.clone())

fuel_tank = LegoAssembly("Fuel Tank", AggregationLayer.ASSEMBLY)
fuel_tank.add_component(battery.clone())

wheels = LegoAssembly("Wheels", AggregationLayer.ASSEMBLY)
for _ in range(4):
    wheels.add_component(wheel.clone())

car.add_assembly(chassis)
car.add_assembly(engine)
car.add_assembly(fuel_tank)
car.add_assembly(wheels)

print_assembly_tree(car)

# Dump to json file
with open("car.json", "w") as fp:
    json.dump(car.to_dict(), fp, cls=KPIEncoder)

pprint.pprint(car.to_dict())
pass