Browse Source

fix: Normalize header handling for all HeaderInit types (#203)

AlexanderSlaa 5 months ago
parent
commit
5007ce7909
2 changed files with 113 additions and 4 deletions
  1. 31 4
      src/utils.ts
  2. 82 0
      test/utils.test.ts

+ 31 - 4
src/utils.ts

@@ -100,6 +100,34 @@ function getPlatform(): string {
   return '' // unknown
 }
 
+/**
+ * Normalizes headers into a plain object format.
+ * This function handles various types of HeaderInit objects such as Headers, arrays of key-value pairs,
+ * and plain objects, converting them all into an object structure.
+ *
+ * @param {HeadersInit|undefined} headers - The headers to normalize. Can be one of the following:
+ *   - A `Headers` object from the Fetch API.
+ *   - A plain object with key-value pairs representing headers.
+ *   - An array of key-value pairs representing headers.
+ * @returns {Record<string,string>} - A plain object representing the normalized headers.
+ */
+function normalizeHeaders(headers?: HeadersInit | undefined): Record<string,string> {
+  if (headers instanceof Headers) {
+      // If headers are an instance of Headers, convert it to an object
+      const obj: Record<string, string> = {};
+        headers.forEach((value, key) => {
+          obj[key] = value;
+        });
+        return obj;
+  } else if (Array.isArray(headers)) {
+      // If headers are in array format, convert them to an object
+      return Object.fromEntries(headers);
+  } else {
+      // Otherwise assume it's already a plain object
+      return headers || {};
+  }
+}
+
 /**
  * A wrapper around fetch that adds default headers.
  * @param fetch {Fetch} - The fetch function to use
@@ -118,10 +146,9 @@ const fetchWithHeaders = async (
     'User-Agent': `ollama-js/${version} (${getPlatform()})`,
   } as HeadersInit
 
-  if (!options.headers) {
-    options.headers = {}
-  }
-
+  // Normalizes headers into a plain object format.
+  options.headers = normalizeHeaders(options.headers);
+  
   // Filter out default headers from custom headers
   const customHeaders = Object.fromEntries(
     Object.entries(options.headers).filter(([key]) => !Object.keys(defaultHeaders).some(defaultKey => defaultKey.toLowerCase() === key.toLowerCase()))

+ 82 - 0
test/utils.test.ts

@@ -0,0 +1,82 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest'
+import { get } from '../src/utils'
+
+describe('get Function Header Tests', () => {
+  const mockFetch = vi.fn();
+  const mockResponse = new Response(null, { status: 200 });
+
+  beforeEach(() => {
+    mockFetch.mockReset();
+    mockFetch.mockResolvedValue(mockResponse);
+  });
+
+  const defaultHeaders = {
+    'Content-Type': 'application/json',
+    'Accept': 'application/json',
+    'User-Agent': expect.stringMatching(/ollama-js\/.*/)
+  };
+
+  it('should use default headers when no headers provided', async () => {
+    await get(mockFetch, 'http://example.com');
+    
+    expect(mockFetch).toHaveBeenCalledWith('http://example.com', {
+      headers: expect.objectContaining(defaultHeaders)
+    });
+  });
+
+  it('should handle Headers instance', async () => {
+    const customHeaders = new Headers({
+      'Authorization': 'Bearer token',
+      'X-Custom': 'value'
+    });
+
+    await get(mockFetch, 'http://example.com', { headers: customHeaders });
+
+    expect(mockFetch).toHaveBeenCalledWith('http://example.com', {
+      headers: expect.objectContaining({
+        ...defaultHeaders,
+        'authorization': 'Bearer token',
+        'x-custom': 'value'
+      })
+    });
+  });
+
+  it('should handle plain object headers', async () => {
+    const customHeaders = {
+      'Authorization': 'Bearer token',
+      'X-Custom': 'value'
+    };
+
+    await get(mockFetch, 'http://example.com', { headers: customHeaders });
+
+    expect(mockFetch).toHaveBeenCalledWith('http://example.com', {
+      headers: expect.objectContaining({
+        ...defaultHeaders,
+        'Authorization': 'Bearer token',
+        'X-Custom': 'value'
+      })
+    });
+  });
+
+  it('should not allow custom headers to override default User-Agent', async () => {
+    const customHeaders = {
+      'User-Agent': 'custom-agent'
+    };
+
+    await get(mockFetch, 'http://example.com', { headers: customHeaders });
+
+    expect(mockFetch).toHaveBeenCalledWith('http://example.com', {
+      headers: expect.objectContaining({
+        'User-Agent': expect.stringMatching(/ollama-js\/.*/)
+      })
+    });
+  });
+
+  it('should handle empty headers object', async () => {
+    await get(mockFetch, 'http://example.com', { headers: {} });
+
+    expect(mockFetch).toHaveBeenCalledWith('http://example.com', {
+      headers: expect.objectContaining(defaultHeaders)
+    });
+  });
+});