|
| 1 | +import { afterEach, describe, expect, it } from 'vitest' |
| 2 | +import { |
| 3 | + mkdtempSync, |
| 4 | + readFileSync, |
| 5 | + rmSync, |
| 6 | + writeFileSync, |
| 7 | + existsSync, |
| 8 | +} from 'node:fs' |
| 9 | +import { tmpdir } from 'node:os' |
| 10 | +import { join } from 'node:path' |
| 11 | +import { PassThrough } from 'node:stream' |
| 12 | +import { ensureAnthropicApiKey, type PromptIO } from '../src/prompt-api-key' |
| 13 | + |
| 14 | +interface FakeIO { |
| 15 | + io: PromptIO |
| 16 | + output: () => string |
| 17 | + cwd: string |
| 18 | + exitCalls: Array<number> |
| 19 | + cleanup: () => void |
| 20 | +} |
| 21 | + |
| 22 | +function createFakeIO({ |
| 23 | + inputLines, |
| 24 | + isTTY = true, |
| 25 | +}: { |
| 26 | + inputLines: Array<string> |
| 27 | + isTTY?: boolean |
| 28 | +}): FakeIO { |
| 29 | + const cwd = mkdtempSync(join(tmpdir(), `electric-ax-prompt-`)) |
| 30 | + const input = new PassThrough() |
| 31 | + const output = new PassThrough() |
| 32 | + let outputBuffer = `` |
| 33 | + let promptsSeen = 0 |
| 34 | + let linesFed = 0 |
| 35 | + const queue = [...inputLines] |
| 36 | + |
| 37 | + output.on(`data`, (chunk: Buffer) => { |
| 38 | + const text = chunk.toString(`utf8`) |
| 39 | + outputBuffer += text |
| 40 | + for (const _match of text.matchAll(/: $/gm)) { |
| 41 | + void _match |
| 42 | + promptsSeen += 1 |
| 43 | + } |
| 44 | + while (linesFed < promptsSeen && linesFed < queue.length) { |
| 45 | + const line = queue[linesFed]! |
| 46 | + linesFed += 1 |
| 47 | + setImmediate(() => input.write(`${line}\n`)) |
| 48 | + } |
| 49 | + }) |
| 50 | + |
| 51 | + const exitCalls: Array<number> = [] |
| 52 | + const exit = ((code: number) => { |
| 53 | + exitCalls.push(code) |
| 54 | + throw new Error(`__exit__:${code}`) |
| 55 | + }) as (code: number) => never |
| 56 | + |
| 57 | + return { |
| 58 | + io: { |
| 59 | + input, |
| 60 | + output, |
| 61 | + isTTY, |
| 62 | + cwd, |
| 63 | + exit, |
| 64 | + }, |
| 65 | + output: () => outputBuffer, |
| 66 | + cwd, |
| 67 | + exitCalls, |
| 68 | + cleanup: () => rmSync(cwd, { recursive: true, force: true }), |
| 69 | + } |
| 70 | +} |
| 71 | + |
| 72 | +describe(`ensureAnthropicApiKey`, () => { |
| 73 | + let cleanups: Array<() => void> = [] |
| 74 | + const originalApiKey = process.env.ANTHROPIC_API_KEY |
| 75 | + |
| 76 | + afterEach(() => { |
| 77 | + for (const cleanup of cleanups) cleanup() |
| 78 | + cleanups = [] |
| 79 | + if (originalApiKey === undefined) { |
| 80 | + delete process.env.ANTHROPIC_API_KEY |
| 81 | + } else { |
| 82 | + process.env.ANTHROPIC_API_KEY = originalApiKey |
| 83 | + } |
| 84 | + }) |
| 85 | + |
| 86 | + it(`returns the explicit option key without prompting`, async () => { |
| 87 | + const fake = createFakeIO({ inputLines: [] }) |
| 88 | + cleanups.push(fake.cleanup) |
| 89 | + |
| 90 | + const key = await ensureAnthropicApiKey( |
| 91 | + { anthropicApiKey: `sk-ant-explicit` }, |
| 92 | + fake.io |
| 93 | + ) |
| 94 | + |
| 95 | + expect(key).toBe(`sk-ant-explicit`) |
| 96 | + expect(fake.output()).toBe(``) |
| 97 | + }) |
| 98 | + |
| 99 | + it(`rethrows the original error when stdin is not a TTY`, async () => { |
| 100 | + delete process.env.ANTHROPIC_API_KEY |
| 101 | + const fake = createFakeIO({ inputLines: [], isTTY: false }) |
| 102 | + cleanups.push(fake.cleanup) |
| 103 | + |
| 104 | + await expect(ensureAnthropicApiKey({}, fake.io)).rejects.toThrow( |
| 105 | + /ANTHROPIC_API_KEY/ |
| 106 | + ) |
| 107 | + expect(fake.output()).toBe(``) |
| 108 | + }) |
| 109 | + |
| 110 | + it(`exits cleanly when the user picks manual setup (1)`, async () => { |
| 111 | + delete process.env.ANTHROPIC_API_KEY |
| 112 | + const fake = createFakeIO({ inputLines: [`1`] }) |
| 113 | + cleanups.push(fake.cleanup) |
| 114 | + |
| 115 | + await expect(ensureAnthropicApiKey({}, fake.io)).rejects.toThrow( |
| 116 | + `__exit__:0` |
| 117 | + ) |
| 118 | + |
| 119 | + expect(fake.exitCalls).toEqual([0]) |
| 120 | + expect(fake.output()).toContain(`never leaves your local computer`) |
| 121 | + expect(fake.output()).toContain(`Get a key from`) |
| 122 | + expect(existsSync(join(fake.cwd, `.env`))).toBe(false) |
| 123 | + }) |
| 124 | + |
| 125 | + it(`writes the pasted key to a new .env when the user picks 2`, async () => { |
| 126 | + delete process.env.ANTHROPIC_API_KEY |
| 127 | + const fake = createFakeIO({ inputLines: [`2`, `sk-ant-pasted`] }) |
| 128 | + cleanups.push(fake.cleanup) |
| 129 | + |
| 130 | + const key = await ensureAnthropicApiKey({}, fake.io) |
| 131 | + |
| 132 | + expect(key).toBe(`sk-ant-pasted`) |
| 133 | + expect(process.env.ANTHROPIC_API_KEY).toBe(`sk-ant-pasted`) |
| 134 | + |
| 135 | + const envContent = readFileSync(join(fake.cwd, `.env`), `utf8`) |
| 136 | + expect(envContent).toBe(`ANTHROPIC_API_KEY=sk-ant-pasted\n`) |
| 137 | + }) |
| 138 | + |
| 139 | + it(`appends the key to an existing .env without an ANTHROPIC_API_KEY entry`, async () => { |
| 140 | + delete process.env.ANTHROPIC_API_KEY |
| 141 | + const fake = createFakeIO({ inputLines: [`2`, `sk-ant-new`] }) |
| 142 | + cleanups.push(fake.cleanup) |
| 143 | + |
| 144 | + writeFileSync(join(fake.cwd, `.env`), `OTHER=value\n`, `utf8`) |
| 145 | + |
| 146 | + await ensureAnthropicApiKey({}, fake.io) |
| 147 | + |
| 148 | + const envContent = readFileSync(join(fake.cwd, `.env`), `utf8`) |
| 149 | + expect(envContent).toBe(`OTHER=value\nANTHROPIC_API_KEY=sk-ant-new\n`) |
| 150 | + }) |
| 151 | + |
| 152 | + it(`replaces an existing empty ANTHROPIC_API_KEY line in .env`, async () => { |
| 153 | + delete process.env.ANTHROPIC_API_KEY |
| 154 | + const fake = createFakeIO({ inputLines: [`2`, `sk-ant-fresh`] }) |
| 155 | + cleanups.push(fake.cleanup) |
| 156 | + |
| 157 | + writeFileSync( |
| 158 | + join(fake.cwd, `.env`), |
| 159 | + `OTHER=value\nANTHROPIC_API_KEY=\nANOTHER=foo\n`, |
| 160 | + `utf8` |
| 161 | + ) |
| 162 | + |
| 163 | + await ensureAnthropicApiKey({}, fake.io) |
| 164 | + |
| 165 | + const envContent = readFileSync(join(fake.cwd, `.env`), `utf8`) |
| 166 | + expect(envContent).toBe( |
| 167 | + `OTHER=value\nANTHROPIC_API_KEY=sk-ant-fresh\nANOTHER=foo\n` |
| 168 | + ) |
| 169 | + }) |
| 170 | + |
| 171 | + it(`rejects an empty pasted key`, async () => { |
| 172 | + delete process.env.ANTHROPIC_API_KEY |
| 173 | + const fake = createFakeIO({ inputLines: [`2`, ``] }) |
| 174 | + cleanups.push(fake.cleanup) |
| 175 | + |
| 176 | + await expect(ensureAnthropicApiKey({}, fake.io)).rejects.toThrow( |
| 177 | + /No API key provided/ |
| 178 | + ) |
| 179 | + }) |
| 180 | + |
| 181 | + it(`rejects an unrecognized choice`, async () => { |
| 182 | + delete process.env.ANTHROPIC_API_KEY |
| 183 | + const fake = createFakeIO({ inputLines: [`bogus`] }) |
| 184 | + cleanups.push(fake.cleanup) |
| 185 | + |
| 186 | + await expect(ensureAnthropicApiKey({}, fake.io)).rejects.toThrow( |
| 187 | + /Unrecognized choice/ |
| 188 | + ) |
| 189 | + }) |
| 190 | +}) |
0 commit comments