tor-browser

The Tor Browser
git clone https://git.dasho.dev/tor-browser.git
Log | Files | Refs | README | LICENSE

test_Utils.js (3663B)

      1 /* This Source Code Form is subject to the terms of the Mozilla Public
      2 * License, v. 2.0. If a copy of the MPL was not distributed with this
      3 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
      4 
      5 const { MODEL_FEATURES, openAIEngine, renderPrompt } =
      6  ChromeUtils.importESModule(
      7    "moz-src:///browser/components/aiwindow/models/Utils.sys.mjs"
      8  );
      9 
     10 const { sinon } = ChromeUtils.importESModule(
     11  "resource://testing-common/Sinon.sys.mjs"
     12 );
     13 
     14 /**
     15 * Constants for preference keys and test values
     16 */
     17 const PREF_API_KEY = "browser.aiwindow.apiKey";
     18 const PREF_ENDPOINT = "browser.aiwindow.endpoint";
     19 const PREF_MODEL = "browser.aiwindow.model";
     20 const PREF_EXTRA_HEADERS = "browser.aiwindow.extraHeaders";
     21 
     22 const API_KEY = "fake-key";
     23 const ENDPOINT = "https://api.fake-endpoint.com/v1";
     24 const MODEL = "fake-model";
     25 const EXTRA_HEADERS = '{"x-fastly-request": "fake-key"}';
     26 
     27 /**
     28 * Cleans up preferences after testing
     29 */
     30 registerCleanupFunction(() => {
     31  for (let pref of [PREF_API_KEY, PREF_ENDPOINT, PREF_MODEL]) {
     32    if (Services.prefs.prefHasUserValue(pref)) {
     33      Services.prefs.clearUserPref(pref);
     34    }
     35  }
     36 });
     37 
     38 /**
     39 * Tests the creation of an OpenAI engine instance
     40 */
     41 add_task(async function test_createOpenAIEngine_with_chat_feature() {
     42  Services.prefs.setStringPref(PREF_API_KEY, API_KEY);
     43  Services.prefs.setStringPref(PREF_ENDPOINT, ENDPOINT);
     44  Services.prefs.setStringPref(PREF_MODEL, MODEL);
     45  Services.prefs.setStringPref(PREF_EXTRA_HEADERS, EXTRA_HEADERS);
     46 
     47  const sb = sinon.createSandbox();
     48  try {
     49    // Take engine to stub out actual engine creation
     50    const fakeEngine = {
     51      runWithGenerator() {
     52        throw new Error("not used");
     53      },
     54    };
     55 
     56    const stub = sb.stub(openAIEngine, "_createEngine").resolves(fakeEngine);
     57    const engine = await openAIEngine.build(MODEL_FEATURES.CHAT);
     58    Assert.strictEqual(
     59      engine.engineInstance,
     60      fakeEngine,
     61      "Should return engine from _createEngine"
     62    );
     63    Assert.ok(stub.calledOnce, "_createEngine should be called once");
     64 
     65    // Test preferences were read correctly
     66    const opts = stub.firstCall.args[0];
     67    Assert.equal(opts.apiKey, API_KEY, "apiKey should come from pref");
     68    Assert.equal(opts.backend, "openai", "backend should be openai");
     69    Assert.equal(opts.baseURL, ENDPOINT, "baseURL should come from pref");
     70    Assert.equal(
     71      opts.engineId,
     72      "smart-openai",
     73      "engineId should be smart-openai"
     74    );
     75    Assert.ok(opts.modelId, "modelId should be set");
     76    Assert.equal(opts.modelRevision, "main", "modelRevision should be main");
     77    Assert.equal(
     78      opts.taskName,
     79      "text-generation",
     80      "taskName should be text-generation"
     81    );
     82    Assert.equal(opts.serviceType, "ai", "serviceType should be ai");
     83    Assert.deepEqual(
     84      opts.extraHeaders,
     85      JSON.parse(EXTRA_HEADERS),
     86      "extraHeaders should come from pref"
     87    );
     88  } finally {
     89    sb.restore();
     90  }
     91 });
     92 
     93 /**
     94 * Tests rendering a prompt from a file with placeholder string replacements
     95 */
     96 add_task(async function test_renderPrompt() {
     97  // Render the test prompt with replacements
     98  const test_prompt = `
     99 This is a test prompt.
    100 {testToReplace1}
    101 
    102 This is more content. {testToReplace2}
    103 
    104 {testToReplace3} Here's the last line.`.trim();
    105  const promptContent = await renderPrompt(test_prompt, {
    106    testToReplace1: "replaced1",
    107    testToReplace2: "replaced2",
    108    testToReplace3: "replaced3",
    109  });
    110 
    111  Assert.equal(
    112    promptContent,
    113    "This is a test prompt.\nreplaced1\n\nThis is more content. replaced2\n\nreplaced3 Here's the last line.",
    114    "Should render the prompt correctly with provided replacement strings"
    115  );
    116 });