diff --git a/calm-hub-ui/src/components/logout-button/LogoutButton.test.tsx b/calm-hub-ui/src/components/logout-button/LogoutButton.test.tsx
new file mode 100644
index 000000000..e226a1baa
--- /dev/null
+++ b/calm-hub-ui/src/components/logout-button/LogoutButton.test.tsx
@@ -0,0 +1,50 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { render, screen, fireEvent } from '@testing-library/react';
+import { LogoutButton } from './LogoutButton.js';
+import { authService } from '../../authService.js';
+
+vi.mock('../../authService.js', () => ({
+    authService: {
+        logout: vi.fn(),
+    },
+}));
+
+describe('LogoutButton', () => {
+    beforeEach(() => {
+        vi.clearAllMocks();
+    });
+
+    it('renders a logout button', () => {
+        render(<LogoutButton />);
+        const button = screen.getByRole('button', { name: /logout/i });
+        expect(button).toBeInTheDocument();
+    });
+
+    it('has correct styling', () => {
+        render(<LogoutButton />);
+        const button = screen.getByRole('button', { name: /logout/i });
+        expect(button).toHaveStyle({
+            position: 'absolute',
+            top: '10px',
+            right: '10px',
+        });
+    });
+
+    it('calls authService.logout when clicked', async () => {
+        render(<LogoutButton />);
+        const button = screen.getByRole('button', { name: /logout/i });
+
+        fireEvent.click(button);
+
+        await new Promise(resolve => setTimeout(resolve, 0));
+
+        expect(authService.logout).toHaveBeenCalled();
+    });
+
+    it('button text is "Logout"', () => {
+        render(<LogoutButton />);
+        const button = screen.getByRole('button');
+        expect(button.textContent).toBe('Logout');
+    });
+});
+
diff --git a/calm-hub-ui/src/components/logout-button/LogoutButton.tsx b/calm-hub-ui/src/components/logout-button/LogoutButton.tsx
new file mode 100644
index 000000000..bf932c495
--- /dev/null
+++ b/calm-hub-ui/src/components/logout-button/LogoutButton.tsx
@@ -0,0 +1,14 @@
+import React from 'react';
+import { authService } from '../../authService.js';
+
+export const LogoutButton: React.FC = () => {
+    const handleLogout = async () => {
+        await authService.logout();
+    };
+
+    return (
+        <button onClick={handleLogout} style={{ position: 'absolute', top: 10, right: 10 }}>
+            Logout
+        </button>
+    );
+};
diff --git a/calm-hub-ui/src/index.tsx b/calm-hub-ui/src/index.tsx
index d498b6c47..e574d9e85 100644
--- a/calm-hub-ui/src/index.tsx
+++ b/calm-hub-ui/src/index.tsx
@@ -2,23 +2,12 @@ import './index.css';
 import React from 'react';
 import ReactDOM from 'react-dom/client';
 import ProtectedRoute from './ProtectedRoute.js';
-import { isAuthServiceEnabled, authService } from './authService.js';
+import { isAuthServiceEnabled } from './authService.js';
 import App from './App.js';
+import { LogoutButton } from './components/logout-button/LogoutButton.js';
 
 const root = ReactDOM.createRoot(document.getElementById('root') as HTMLElement);
 
-const LogoutButton: React.FC = () => {
-    const handleLogout = async () => {
-        await authService.logout();
-    };
-
-    return (
-        <button onClick={handleLogout} style={{ position: 'absolute', top: 10, right: 10 }}>
-            Logout
-        </button>
-    );
-};
-
 const isAuthenticationEnabled = isAuthServiceEnabled();
 
 root.render(
diff --git a/calm-hub-ui/src/visualizer/components/reactflow/edge-components/EdgeBadge.test.tsx b/calm-hub-ui/src/visualizer/components/reactflow/edge-components/EdgeBadge.test.tsx
index 00ae95083..253d2c42a 100644
--- a/calm-hub-ui/src/visualizer/components/reactflow/edge-components/EdgeBadge.test.tsx
+++ b/calm-hub-ui/src/visualizer/components/reactflow/edge-components/EdgeBadge.test.tsx
@@ -1,7 +1,8 @@
 import { describe, it, expect, vi } from 'vitest';
 import { render, fireEvent } from '@testing-library/react';
-import { EdgeBadge, getBadgeStyle } from './EdgeBadge';
-import { THEME } from '../theme';
+import { EdgeBadge } from './EdgeBadge.js';
+import { getBadgeStyle } from '../utils/edgeBadge.utils.js';
+import { THEME } from '../theme.js';
 
 describe('EdgeBadge', () => {
     const mockOnMouseEnter = vi.fn();
diff --git a/calm-hub-ui/src/visualizer/components/reactflow/edge-components/EdgeBadge.tsx b/calm-hub-ui/src/visualizer/components/reactflow/edge-components/EdgeBadge.tsx
index b0f866a91..50cbcb019 100644
--- a/calm-hub-ui/src/visualizer/components/reactflow/edge-components/EdgeBadge.tsx
+++ b/calm-hub-ui/src/visualizer/components/reactflow/edge-components/EdgeBadge.tsx
@@ -1,6 +1,5 @@
 import { Info, Shield, ArrowRight } from 'lucide-react';
-import { THEME } from '../theme.js';
-import type { EdgeBadgeStyle, EdgeBadgeProps } from '../../../contracts/contracts.js';
+import type { EdgeBadgeProps } from '../../../contracts/contracts.js';
 
 export function EdgeBadge({
     hasFlowInfo,
@@ -36,25 +35,3 @@ export function EdgeBadge({
         </div>
     );
 }
-
-export function getBadgeStyle(hasFlowInfo: boolean, hasAIGF: boolean): EdgeBadgeStyle {
-    if (hasFlowInfo) {
-        return {
-            background: `${THEME.colors.accent}20`,
-            border: THEME.colors.accent,
-            iconColor: THEME.colors.accent,
-        };
-    }
-    if (hasAIGF) {
-        return {
-            background: `${THEME.colors.success}20`,
-            border: THEME.colors.success,
-            iconColor: THEME.colors.success,
-        };
-    }
-    return {
-        background: `${THEME.colors.muted}20`,
-        border: THEME.colors.muted,
-        iconColor: THEME.colors.muted,
-    };
-}
diff --git a/calm-hub-ui/src/visualizer/components/reactflow/edge-components/index.ts b/calm-hub-ui/src/visualizer/components/reactflow/edge-components/index.ts
index 429165199..29a5eeae1 100644
--- a/calm-hub-ui/src/visualizer/components/reactflow/edge-components/index.ts
+++ b/calm-hub-ui/src/visualizer/components/reactflow/edge-components/index.ts
@@ -1,4 +1,5 @@
-export { EdgeBadge, getBadgeStyle } from './EdgeBadge.js';
+export { EdgeBadge } from './EdgeBadge.js';
+export { getBadgeStyle } from '../utils/edgeBadge.utils.js';
 export { EdgeTooltip } from './EdgeTooltip.js';
 export type {
     EdgeBadgeProps,
diff --git a/calm-hub-ui/src/visualizer/components/reactflow/utils/edgeBadge.utils.test.ts b/calm-hub-ui/src/visualizer/components/reactflow/utils/edgeBadge.utils.test.ts
new file mode 100644
index 000000000..a66e6a65e
--- /dev/null
+++ b/calm-hub-ui/src/visualizer/components/reactflow/utils/edgeBadge.utils.test.ts
@@ -0,0 +1,48 @@
+import { describe, it, expect } from 'vitest';
+import { getBadgeStyle } from './edgeBadge.utils.js';
+import { THEME } from '../theme.js';
+
+describe('edgeBadge.utils', () => {
+    describe('getBadgeStyle', () => {
+        it('returns accent colors when hasFlowInfo is true', () => {
+            const result = getBadgeStyle(true, false);
+            expect(result).toEqual({
+                background: `${THEME.colors.accent}20`,
+                border: THEME.colors.accent,
+                iconColor: THEME.colors.accent,
+            });
+        });
+
+        it('returns success colors when hasAIGF is true', () => {
+            const result = getBadgeStyle(false, true);
+            expect(result).toEqual({
+                background: `${THEME.colors.success}20`,
+                border: THEME.colors.success,
+                iconColor: THEME.colors.success,
+            });
+        });
+
+        it('returns muted colors when both flags are false', () => {
+            const result = getBadgeStyle(false, false);
+            expect(result).toEqual({
+                background: `${THEME.colors.muted}20`,
+                border: THEME.colors.muted,
+                iconColor: THEME.colors.muted,
+            });
+        });
+
+        it('prioritizes hasFlowInfo over hasAIGF when both are true', () => {
+            const result = getBadgeStyle(true, true);
+            expect(result).toEqual({
+                background: `${THEME.colors.accent}20`,
+                border: THEME.colors.accent,
+                iconColor: THEME.colors.accent,
+            });
+        });
+
+        it('has the correct alpha value for background colors', () => {
+            const result = getBadgeStyle(true, false);
+            expect(result.background).toMatch(/^#[0-9a-f]+20$/i);
+        });
+    });
+});
diff --git a/calm-hub-ui/src/visualizer/components/reactflow/utils/edgeBadge.utils.ts b/calm-hub-ui/src/visualizer/components/reactflow/utils/edgeBadge.utils.ts
new file mode 100644
index 000000000..f200b8e4c
--- /dev/null
+++ b/calm-hub-ui/src/visualizer/components/reactflow/utils/edgeBadge.utils.ts
@@ -0,0 +1,24 @@
+import { THEME } from '../theme.js';
+import type { EdgeBadgeStyle } from '../../../contracts/contracts.js';
+
+export function getBadgeStyle(hasFlowInfo: boolean, hasAIGF: boolean): EdgeBadgeStyle {
+    if (hasFlowInfo) {
+        return {
+            background: `${THEME.colors.accent}20`,
+            border: THEME.colors.accent,
+            iconColor: THEME.colors.accent,
+        };
+    }
+    if (hasAIGF) {
+        return {
+            background: `${THEME.colors.success}20`,
+            border: THEME.colors.success,
+            iconColor: THEME.colors.success,
+        };
+    }
+    return {
+        background: `${THEME.colors.muted}20`,
+        border: THEME.colors.muted,
+        iconColor: THEME.colors.muted,
+    };
+}