Repository: streamlit/streamlit
File: frontend/lib/src/render-tree/visitors/ElementsSetVisitor.test.ts (TypeScript)

Code before:
      const child2 = text("child2")
      const blockNode = block([child1, child2])
      const visitor = new ElementsSetVisitor()

      visitor.visitBlockNode(blockNode)

      expect(visitor.elements.size).toBe(2)
      expect(visitor.elements.has(child1.element)).toBe(true)
      expect(visitor.elements.has(child2.element)).toBe(true)
    })

    it("handles nested block structures", () => {
      const innerElement = text("inner")
      const outerElement = text("outer")
      const innerBlock = block([innerElement])
      const outerBlock = block([outerElement, innerBlock])
      const visitor = new ElementsSetVisitor()

      visitor.visitBlockNode(outerBlock)

      expect(visitor.elements.size).toBe(2)
      expect(visitor.elements.has(innerElement.element)).toBe(true)
      expect(visitor.elements.has(outerElement.element)).toBe(true)
    })

    it("accumulates elements with existing elements", () => {
      const existingElement = text("existing")
      const blockElement = text("from_block")
      const blockNode = block([blockElement])
      const visitor = new ElementsSetVisitor()

      // First add an existing element
      visitor.visitElementNode(existingElement)
      // Then visit the block
      visitor.visitBlockNode(blockNode)

      expect(visitor.elements.size).toBe(2)
      expect(visitor.elements.has(existingElement.element)).toBe(true)
      expect(visitor.elements.has(blockElement.element)).toBe(true)
    })
  })

  describe("collectElements static method", () => {
    it("collects elements from ElementNode", () => {
      const element = text("test")

      const elements = ElementsSetVisitor.collectElements(element)

      expect(elements.size).toBe(1)
      expect(elements.has(element.element)).toBe(true)
    })

Review comment (suggestion):
[suggestion] I think this description could be differentiated more from the previous "accumulates multiple elements" test. Perhaps, "accumulates multiple elements when inside block nodes"? The emphasis on existing elements seems less important to me. 

Code after:
      const child2 = text("child2")
      const blockNode = block([child1, child2])
      const visitor = new ElementsSetVisitor()

      visitor.visitBlockNode(blockNode)

      expect(visitor.elements.size).toBe(2)
      expect(visitor.elements.has(child1.element)).toBe(true)
      expect(visitor.elements.has(child2.element)).toBe(true)
    })

    it("handles nested block structures", () => {
      const innerElement = text("inner")
      const dupeElement = innerElement
      const outerElement = text("outer")
      const innerBlock = block([innerElement])
      const outerBlock = block([outerElement, innerBlock, dupeElement])
      const visitor = new ElementsSetVisitor()

      visitor.visitBlockNode(outerBlock)

      // Ensure the dupe element is removed
      expect(visitor.elements.size).toBe(2)
      expect(visitor.elements.has(innerElement.element)).toBe(true)
      expect(visitor.elements.has(outerElement.element)).toBe(true)
    })

    it("accumulates multiple elements when inside block nodes", () => {
      const existingElement = text("existing")
      const blockElement = text("from_block")
      const blockNode = block([blockElement])
      const visitor = new ElementsSetVisitor()

      // First add an existing element
      visitor.visitElementNode(existingElement)
      // Then visit the block
      visitor.visitBlockNode(blockNode)

      expect(visitor.elements.size).toBe(2)
      expect(visitor.elements.has(existingElement.element)).toBe(true)
      expect(visitor.elements.has(blockElement.element)).toBe(true)
    })
  })

  describe("collectElements static method", () => {
    it("collects elements from ElementNode", () => {
      const element = text("test")
      const elements = ElementsSetVisitor.collectElements(element)

      expect(elements.size).toBe(1)
      expect(elements.has(element.element)).toBe(true)