From f2d188f24c7c0a5697956b97a98431eb8d087f1f Mon Sep 17 00:00:00 2001
From: Tyler Dixon <tylerdixon@github.com>
Date: Thu, 12 Dec 2024 09:26:34 -0500
Subject: [PATCH] Add className back to TokenBase (#5420)

* Add className back to TokenBase

* Adding tests and fixing implementation

* Create dirty-wasps-behave.md

---------

Co-authored-by: Marie Lucca <40550942+francinelucca@users.noreply.github.com>
---
 .changeset/dirty-wasps-behave.md                  | 5 +++++
 packages/react/src/Token/AvatarToken.tsx          | 1 +
 packages/react/src/Token/IssueLabelToken.tsx      | 1 +
 packages/react/src/Token/Token.tsx                | 1 +
 packages/react/src/Token/TokenBase.tsx            | 1 +
 packages/react/src/Token/__tests__/Token.test.tsx | 7 +++++++
 6 files changed, 16 insertions(+)
 create mode 100644 .changeset/dirty-wasps-behave.md

diff --git a/.changeset/dirty-wasps-behave.md b/.changeset/dirty-wasps-behave.md
new file mode 100644
index 00000000000..f0d97a4ff78
--- /dev/null
+++ b/.changeset/dirty-wasps-behave.md
@@ -0,0 +1,5 @@
+---
+"@primer/react": patch
+---
+
+Add className back to TokenBase
diff --git a/packages/react/src/Token/AvatarToken.tsx b/packages/react/src/Token/AvatarToken.tsx
index c6b70005639..b658fb7df30 100644
--- a/packages/react/src/Token/AvatarToken.tsx
+++ b/packages/react/src/Token/AvatarToken.tsx
@@ -70,6 +70,7 @@ const AvatarToken = forwardRef(({avatarSrc, id, size = defaultTokenSize, classNa
       sx={{
         paddingLeft: get('space.1'),
       }}
+      className={className}
       {...rest}
       ref={forwardedRef}
     />
diff --git a/packages/react/src/Token/IssueLabelToken.tsx b/packages/react/src/Token/IssueLabelToken.tsx
index a8b1f409d67..830eb0f0e12 100644
--- a/packages/react/src/Token/IssueLabelToken.tsx
+++ b/packages/react/src/Token/IssueLabelToken.tsx
@@ -180,6 +180,7 @@ const IssueLabelToken = forwardRef((props, forwardedRef) => {
       text={text}
       size={size}
       sx={labelStyles}
+      className={className}
       {...(!hasMultipleActionTargets ? interactiveTokenProps : {})}
       {...rest}
       ref={forwardedRef}
diff --git a/packages/react/src/Token/Token.tsx b/packages/react/src/Token/Token.tsx
index c13e4807920..49a6183f52a 100644
--- a/packages/react/src/Token/Token.tsx
+++ b/packages/react/src/Token/Token.tsx
@@ -136,6 +136,7 @@ const Token = forwardRef((props, forwardedRef) => {
       text={text}
       size={size}
       sx={mergedSx}
+      className={className}
       {...(!hasMultipleActionTargets ? interactiveTokenProps : {})}
       {...rest}
       ref={forwardedRef}
diff --git a/packages/react/src/Token/TokenBase.tsx b/packages/react/src/Token/TokenBase.tsx
index 89b3b28219a..74bee51336c 100644
--- a/packages/react/src/Token/TokenBase.tsx
+++ b/packages/react/src/Token/TokenBase.tsx
@@ -173,6 +173,7 @@ const TokenBase = React.forwardRef<HTMLButtonElement | HTMLAnchorElement | HTMLS
             onRemove()
           }
         }}
+        className={className}
         id={id?.toString()}
         size={size}
         {...rest}
diff --git a/packages/react/src/Token/__tests__/Token.test.tsx b/packages/react/src/Token/__tests__/Token.test.tsx
index 45274a3d937..ad74c75f646 100644
--- a/packages/react/src/Token/__tests__/Token.test.tsx
+++ b/packages/react/src/Token/__tests__/Token.test.tsx
@@ -77,6 +77,13 @@ const testTokenComponent = (Component: React.ComponentType<React.PropsWithChildr
     expect(onRemoveMock).toHaveBeenCalled()
   })
 
+  it('adds className to rendered component', () => {
+    const {getByText} = HTMLRender(<Component text="token" className="testing-class" />)
+    const domNode = getByText('token')
+
+    expect(domNode.parentElement).toHaveClass('testing-class')
+  })
+
   it('should have no axe violations', async () => {
     const {container} = HTMLRender(<Component text="token" />)
     const results = await axe.run(container)