Fix the testing type badges

src/components/badge/index.tsx

@@ -2,15 +2,15 @@ import React from "react";
 import s from "./style.module.css";
 import { BadgeProps } from "./types";
 
-function classNames(...classes) {
+function classNames(...classes: string[]) {
   return classes.filter(Boolean).join(" ");
 }
 
 export default function Badge({ type, path, children }: BadgeProps) {
   return (
     <>
       {path && (
-        <a href={path}>
+        <a href={path} style={{ borderBottom: 'none' }}>
           <div className={classNames(`${s.badge}`, `${s[type]}`)}>
             {children}
           </div>

src/components/product-heading/index.tsx

@@ -1,19 +1,24 @@
 import Icon from '@cypress-design/react-icon'
 import Badge from "@site/src/components/badge"
 import s from './style.module.css'
+import {useDoc} from '@docusaurus/theme-common/internal';
+import E2EOnlyBadge from "@site/src/components/e2e-only-badge";
+import ComponentOnlyBadge from "@site/src/components/component-only-badge";
 
 import React from 'react';
 
 // Define the types for the props
 interface ProductHeadingProps {
     product: 'app' | 'cloud' | 'accessibility' | 'ui-coverage'
     plan?: 'team' | 'business' | 'enterprise'
+    badge?: React.ReactNode
 }
 
 // Build the Button component with the specified props
-const ProductHeading: React.FC<ProductHeadingProps> = ({ 
+const DocProductHeading: React.FC<ProductHeadingProps> = ({ 
     product, // The product to display
     plan, // The plan to display for Cloud product
+    badge, // The badge to display
 }) => {
     const productName = product === 'ui-coverage' ? 'UI Coverage' : product === 'accessibility' ? 'Cypress Accessibility' : product === 'cloud' ? 'Cypress Cloud' : 'Cypress App'
     const iconName = product === 'ui-coverage' ? 'technology-ui-coverage' : product === 'accessibility' ? 'cypress-accessibility-outline' : 'technology-cypress'
@@ -26,15 +31,14 @@ const ProductHeading: React.FC<ProductHeadingProps> = ({
     }
 
     return (
-    <div className={s.productHeading}>
+    <div className={s.productHeading} style={{display: 'flex', alignItems: 'center', gap: '0.5rem'}}>
         <Icon 
             name={iconName} 
             className={s.productHeadingIcon}
         />
         <span className={s.productHeadingText}>
           {productName}
         </span>
-
         <a 
           href={`https://www.cypress.io/${linkPath}?utm_source=docs&utm_medium=product-heading-${product}&utm_content=${badgeContent}`}
           target="_blank"
@@ -43,9 +47,18 @@ const ProductHeading: React.FC<ProductHeadingProps> = ({
           >
             { product !== 'app' &&  <Badge type="success">{badgeContent}</Badge>}
         </a>
-
+        <span className={s.productHeadingBadge}>{badge}</span>
     </div>
     )
 }
 
-export default ProductHeading
+const ProductHeading: React.FC<Omit<ProductHeadingProps, 'badge'>> = (props) => {
+  const { frontMatter } = useDoc();
+  const e2eSpecific = (frontMatter as any).e2eSpecific;
+  const componentSpecific = (frontMatter as any).componentSpecific;
+  const testTypePill = (e2eSpecific && <E2EOnlyBadge />) || (componentSpecific && <ComponentOnlyBadge />);
+  return <DocProductHeading {...props} badge={testTypePill} />;
+};
+
+export default ProductHeading;
+export { ProductHeading, DocProductHeading };

src/components/product-heading/style.module.css

@@ -16,4 +16,10 @@
 .productHeadingText {
   font-size: 1.2rem;
   margin-right: 0.5rem;
+}
+
+.productHeadingBadge {
+  margin-left: auto;
+  display: flex;
+  align-items: center;
 }

src/theme/DocItem/Content/index.tsx

@@ -28,10 +28,7 @@ function useSyntheticTitle() {
   }
   return metadata.title;
 }
-export default function DocItemContent({children}) {
-
-const { frontMatter: {e2eSpecific, componentSpecific} } = useDoc();
-const testTypePill = e2eSpecific && <E2EOnlyBadge /> || componentSpecific && <ComponentOnlyBadge />
+export default function DocItemContent({children}: {children: React.ReactNode}) {
 
 const syntheticTitle = useSyntheticTitle();
 return (
@@ -41,7 +38,6 @@ return (
           <div className={s.mainContentHeader}>
             <div className={s.headerWrapper}>
               <Heading as="h1">{syntheticTitle}</Heading>
-              {testTypePill}
             </div>
           </div>
         </header>