From 02f661a103a858aec147d2232039aa14e50f2e54 Mon Sep 17 00:00:00 2001
From: yari-dewalt <yari.dewalt@gmail.com>
Date: Thu, 23 Jan 2025 12:04:04 -0800
Subject: [PATCH] Update and add tests

---
 .../parser/requirementDiagram.spec.js         | 249 +++++++++++++++++-
 .../requirement/requirementDb.spec.ts         |  95 +++++++
 2 files changed, 343 insertions(+), 1 deletion(-)
 create mode 100644 packages/mermaid/src/diagrams/requirement/requirementDb.spec.ts

diff --git a/packages/mermaid/src/diagrams/requirement/parser/requirementDiagram.spec.js b/packages/mermaid/src/diagrams/requirement/parser/requirementDiagram.spec.js
index dce62c6ed..743543f28 100644
--- a/packages/mermaid/src/diagrams/requirement/parser/requirementDiagram.spec.js
+++ b/packages/mermaid/src/diagrams/requirement/parser/requirementDiagram.spec.js
@@ -37,7 +37,7 @@ describe('when parsing requirement diagram it...', function () {
 
     let foundReq = requirementDb.getRequirements().get(expectedName);
     expect(foundReq).toBeDefined();
-    expect(foundReq.id).toBe(expectedId);
+    expect(foundReq.requirementId).toBe(expectedId);
     expect(foundReq.text).toBe(expectedText);
 
     expect(requirementDb.getElements().size).toBe(0);
@@ -599,4 +599,251 @@ line 2`;
       expect(reqDiagram.parser.yy.getElements().size).toBe(1);
     });
   }
+
+  it('will accept styling a requirement', function () {
+    const expectedName = 'test_req';
+
+    let lines = [
+      `requirementDiagram`,
+      ``,
+      `requirement ${expectedName} {`,
+      `}`,
+      `style ${expectedName} fill:#f9f,stroke:#333,stroke-width:4px`,
+      ``,
+    ];
+    let doc = lines.join('\n');
+
+    reqDiagram.parser.parse(doc);
+
+    let foundReq = requirementDb.getRequirements().get(expectedName);
+    const styles = foundReq.cssStyles;
+    expect(styles).toEqual(['fill:#f9f', 'stroke:#333', 'stroke-width:4px']);
+  });
+
+  it('will accept styling an element', function () {
+    const expectedName = 'test_element';
+
+    let lines = [
+      `requirementDiagram`,
+      ``,
+      `element ${expectedName} {`,
+      `}`,
+      `style ${expectedName} fill:#f9f,stroke:#333,stroke-width:4px`,
+      ``,
+    ];
+    let doc = lines.join('\n');
+
+    reqDiagram.parser.parse(doc);
+
+    let foundElement = requirementDb.getElements().get(expectedName);
+    const styles = foundElement.cssStyles;
+    expect(styles).toEqual(['fill:#f9f', 'stroke:#333', 'stroke-width:4px']);
+  });
+
+  it('will accept styling multiple things at once', function () {
+    const expectedRequirementName = 'test_requirement';
+    const expectedElementName = 'test_element';
+
+    let lines = [
+      `requirementDiagram`,
+      ``,
+      `requirement ${expectedRequirementName} {`,
+      `}`,
+      `element ${expectedElementName} {`,
+      `}`,
+      `style ${expectedRequirementName},${expectedElementName} fill:#f9f,stroke:#333,stroke-width:4px`,
+      ``,
+    ];
+    let doc = lines.join('\n');
+
+    reqDiagram.parser.parse(doc);
+
+    let foundRequirement = requirementDb.getRequirements().get(expectedRequirementName);
+    const requirementStyles = foundRequirement.cssStyles;
+    expect(requirementStyles).toEqual(['fill:#f9f', 'stroke:#333', 'stroke-width:4px']);
+    let foundElement = requirementDb.getElements().get(expectedElementName);
+    const elementStyles = foundElement.cssStyles;
+    expect(elementStyles).toEqual(['fill:#f9f', 'stroke:#333', 'stroke-width:4px']);
+  });
+
+  it('will accept defining a class', function () {
+    const expectedName = 'myClass';
+
+    let lines = [
+      `requirementDiagram`,
+      ``,
+      `classDef ${expectedName} fill:#f9f,stroke:#333,stroke-width:4px`,
+      ``,
+    ];
+    let doc = lines.join('\n');
+
+    reqDiagram.parser.parse(doc);
+
+    let foundClass = requirementDb.getClasses().get(expectedName);
+    expect(foundClass).toEqual({
+      id: 'myClass',
+      styles: ['fill:#f9f', 'stroke:#333', 'stroke-width:4px'],
+      textStyles: [],
+    });
+  });
+
+  it('will accept defining multiple classes at once', function () {
+    const firstName = 'firstClass';
+    const secondName = 'secondClass';
+
+    let lines = [
+      `requirementDiagram`,
+      ``,
+      `classDef ${firstName},${secondName} fill:#f9f,stroke:#333,stroke-width:4px`,
+      ``,
+    ];
+    let doc = lines.join('\n');
+
+    reqDiagram.parser.parse(doc);
+
+    let firstClass = requirementDb.getClasses().get(firstName);
+    expect(firstClass).toEqual({
+      id: 'firstClass',
+      styles: ['fill:#f9f', 'stroke:#333', 'stroke-width:4px'],
+      textStyles: [],
+    });
+    let secondClass = requirementDb.getClasses().get(secondName);
+    expect(secondClass).toEqual({
+      id: 'secondClass',
+      styles: ['fill:#f9f', 'stroke:#333', 'stroke-width:4px'],
+      textStyles: [],
+    });
+  });
+
+  it('will accept assigning a class via the class statement', function () {
+    const requirementName = 'myReq';
+    const className = 'myClass';
+
+    let lines = [
+      `requirementDiagram`,
+      ``,
+      `requirement ${requirementName} {`,
+      `}`,
+      `classDef ${className} fill:#f9f,stroke:#333,stroke-width:4px`,
+      `class ${requirementName} ${className}`,
+      ``,
+    ];
+    let doc = lines.join('\n');
+
+    reqDiagram.parser.parse(doc);
+
+    let foundRequirement = requirementDb.getRequirements().get(requirementName);
+    expect(foundRequirement.classes).toEqual(['default', className]);
+  });
+
+  it('will accept assigning multiple classes to multiple things via the class statement', function () {
+    const requirementName = 'req';
+    const elementName = 'elem';
+    const firstClassName = 'class1';
+    const secondClassName = 'class2';
+
+    let lines = [
+      `requirementDiagram`,
+      ``,
+      `requirement ${requirementName} {`,
+      `}`,
+      `element ${elementName} {`,
+      `}`,
+      `classDef ${firstClassName},${secondClassName} fill:#f9f,stroke:#333,stroke-width:4px`,
+      `class ${requirementName},${elementName} ${firstClassName},${secondClassName}`,
+      ``,
+    ];
+    let doc = lines.join('\n');
+
+    reqDiagram.parser.parse(doc);
+
+    let requirement = requirementDb.getRequirements().get(requirementName);
+    expect(requirement.classes).toEqual(['default', firstClassName, secondClassName]);
+    let element = requirementDb.getElements().get(elementName);
+    expect(element.classes).toEqual(['default', firstClassName, secondClassName]);
+  });
+
+  it('will accept assigning a class via the shorthand syntax', function () {
+    const requirementName = 'myReq';
+    const className = 'myClass';
+
+    let lines = [
+      `requirementDiagram`,
+      ``,
+      `requirement ${requirementName} {`,
+      `}`,
+      `classDef ${className} fill:#f9f,stroke:#333,stroke-width:4px`,
+      `${requirementName}:::${className}`,
+      ``,
+    ];
+    let doc = lines.join('\n');
+
+    reqDiagram.parser.parse(doc);
+
+    let foundRequirement = requirementDb.getRequirements().get(requirementName);
+    expect(foundRequirement.classes).toEqual(['default', className]);
+  });
+
+  it('will accept assigning multiple classes via the shorthand syntax', function () {
+    const requirementName = 'myReq';
+    const firstClassName = 'class1';
+    const secondClassName = 'class2';
+
+    let lines = [
+      `requirementDiagram`,
+      ``,
+      `requirement ${requirementName} {`,
+      `}`,
+      `classDef ${firstClassName} fill:#f9f,stroke:#333,stroke-width:4px`,
+      `classDef ${secondClassName} color:blue`,
+      `${requirementName}:::${firstClassName},${secondClassName}`,
+      ``,
+    ];
+    let doc = lines.join('\n');
+
+    reqDiagram.parser.parse(doc);
+
+    let foundRequirement = requirementDb.getRequirements().get(requirementName);
+    expect(foundRequirement.classes).toEqual(['default', firstClassName, secondClassName]);
+  });
+
+  it('will accept assigning a class or multiple via the shorthand syntax when defining a requirement or element', function () {
+    const requirementName = 'myReq';
+    const elementName = 'myElem';
+    const firstClassName = 'class1';
+    const secondClassName = 'class2';
+
+    let lines = [
+      `requirementDiagram`,
+      ``,
+      `requirement ${requirementName}:::${firstClassName} {`,
+      `}`,
+      `element ${elementName}:::${firstClassName},${secondClassName} {`,
+      `}`,
+      ``,
+      `classDef ${firstClassName} fill:#f9f,stroke:#333,stroke-width:4px`,
+      `classDef ${secondClassName} color:blue`,
+      ``,
+    ];
+    let doc = lines.join('\n');
+
+    reqDiagram.parser.parse(doc);
+
+    let foundRequirement = requirementDb.getRequirements().get(requirementName);
+    expect(foundRequirement.classes).toEqual(['default', firstClassName]);
+    let foundElement = requirementDb.getElements().get(elementName);
+    expect(foundElement.classes).toEqual(['default', firstClassName, secondClassName]);
+  });
+
+  describe('will parse direction statements and', () => {
+    test.each(['TB', 'BT', 'LR', 'RL'])('will accept direction %s', (directionVal) => {
+      const lines = ['requirementDiagram', '', `direction ${directionVal}`, ''];
+      const doc = lines.join('\n');
+
+      reqDiagram.parser.parse(doc);
+
+      const direction = requirementDb.getDirection();
+      expect(direction).toBe(directionVal);
+    });
+  });
 });
diff --git a/packages/mermaid/src/diagrams/requirement/requirementDb.spec.ts b/packages/mermaid/src/diagrams/requirement/requirementDb.spec.ts
new file mode 100644
index 000000000..da6684ca1
--- /dev/null
+++ b/packages/mermaid/src/diagrams/requirement/requirementDb.spec.ts
@@ -0,0 +1,95 @@
+import requirementDb from './requirementDb.js';
+import { describe, it, expect } from 'vitest';
+import type { RelationshipType } from './types.js';
+
+describe('requirementDb', () => {
+  beforeEach(() => {
+    requirementDb.clear();
+  });
+
+  it('should add a requirement', () => {
+    requirementDb.addRequirement('requirement', 'Requirement');
+    const requirements = requirementDb.getRequirements();
+    expect(requirements.has('requirement')).toBe(true);
+  });
+
+  it('should add an element', () => {
+    requirementDb.addElement('element');
+    const elements = requirementDb.getElements();
+    expect(elements.has('element')).toBe(true);
+  });
+
+  it('should add a relationship', () => {
+    requirementDb.addRelationship('contains' as RelationshipType, 'src', 'dst');
+    const relationships = requirementDb.getRelationships();
+    const relationship = relationships.find(
+      (r) => r.type === 'contains' && r.src === 'src' && r.dst === 'dst'
+    );
+    expect(relationship).toBeDefined();
+  });
+
+  it('should detect single class', () => {
+    requirementDb.defineClass(['a'], ['stroke-width: 8px']);
+    const classes = requirementDb.getClasses();
+
+    expect(classes.has('a')).toBe(true);
+    expect(classes.get('a')?.styles).toEqual(['stroke-width: 8px']);
+  });
+
+  it('should detect many classes', () => {
+    requirementDb.defineClass(['a', 'b'], ['stroke-width: 8px']);
+    const classes = requirementDb.getClasses();
+
+    expect(classes.has('a')).toBe(true);
+    expect(classes.has('b')).toBe(true);
+    expect(classes.get('a')?.styles).toEqual(['stroke-width: 8px']);
+    expect(classes.get('b')?.styles).toEqual(['stroke-width: 8px']);
+  });
+
+  it('should detect direction', () => {
+    requirementDb.setDirection('TB');
+    const direction = requirementDb.getDirection();
+
+    expect(direction).toBe('TB');
+  });
+
+  it('should add styles to a requirement and element', () => {
+    requirementDb.addRequirement('requirement', 'Requirement');
+    requirementDb.setCssStyle(['requirement'], ['color:red']);
+    requirementDb.addElement('element');
+    requirementDb.setCssStyle(['element'], ['stroke-width:4px', 'stroke: yellow']);
+
+    const requirement = requirementDb.getRequirements().get('requirement');
+    const element = requirementDb.getElements().get('element');
+
+    expect(requirement?.cssStyles).toEqual(['color:red']);
+    expect(element?.cssStyles).toEqual(['stroke-width:4px', 'stroke: yellow']);
+  });
+
+  it('should add classes to a requirement and element', () => {
+    requirementDb.addRequirement('requirement', 'Requirement');
+    requirementDb.addElement('element');
+    requirementDb.setClass(['requirement', 'element'], ['myClass']);
+
+    const requirement = requirementDb.getRequirements().get('requirement');
+    const element = requirementDb.getElements().get('element');
+
+    expect(requirement?.classes).toEqual(['default', 'myClass']);
+    expect(element?.classes).toEqual(['default', 'myClass']);
+  });
+
+  it('should add styles to a requirement and element inherited from a class', () => {
+    requirementDb.addRequirement('requirement', 'Requirement');
+    requirementDb.addElement('element');
+    requirementDb.defineClass(['myClass'], ['color:red']);
+    requirementDb.defineClass(['myClass2'], ['stroke-width:4px', 'stroke: yellow']);
+    requirementDb.setClass(['requirement'], ['myClass']);
+    requirementDb.setClass(['element'], ['myClass2']);
+
+    const requirement = requirementDb.getRequirements().get('requirement');
+    const element = requirementDb.getElements().get('element');
+
+    expect(requirement?.cssStyles).toEqual(['color:red']);
+    expect(element?.cssStyles).toEqual(['stroke-width:4px', 'stroke: yellow']);
+  });
+});