fix: charge bots for traffic

piece-retriever/bin/piece-retriever.js

@@ -230,7 +230,6 @@ export default {
             egressBytes,
             cacheMiss: retrievalResult.cacheMiss,
             enforceEgressQuota: env.ENFORCE_EGRESS_QUOTA,
-            isBotTraffic: !!botName,
           })
         })(),
       )

piece-retriever/lib/store.js

@@ -253,13 +253,7 @@ export async function getRetrievalCandidatesAndValidatePayer(
  */
 export async function updateDataSetStats(
   env,
-  {
-    dataSetId,
-    egressBytes,
-    cacheMiss,
-    enforceEgressQuota = false,
-    isBotTraffic = false,
-  },
+  { dataSetId, egressBytes, cacheMiss, enforceEgressQuota = false },
 ) {
   await env.DB.prepare(
     `
@@ -271,7 +265,7 @@ export async function updateDataSetStats(
     .bind(egressBytes, dataSetId)
     .run()
 
-  if (enforceEgressQuota && !isBotTraffic) {
+  if (enforceEgressQuota) {
     await env.DB.prepare(
       `
       UPDATE data_set_egress_quotas

piece-retriever/test/retriever.test.js

@@ -462,6 +462,37 @@ describe('piece-retriever.fetch', () => {
     },
   )
 
+  it('charges bots for egress', async () => {
+    const botToken = Object.keys(botTokens)[0]
+    /** @type {string} */
+    const botName = env.BOT_TOKENS[botToken]
+    console.log({ botToken, botName })
+
+    const mockRetrieveFile = vi.fn().mockResolvedValue({
+      response: new Response('fake'),
+      cacheMiss: true,
+    })
+    const ctx = createExecutionContext()
+    const req = withRequest(defaultPayerAddress, realPieceCid, 'GET', {
+      authorization: `Bearer ${botToken}`,
+    })
+    const res = await worker.fetch(req, env, ctx, {
+      retrieveFile: mockRetrieveFile,
+    })
+    await waitOnExecutionContext(ctx)
+    expect(res.status).toBe(200)
+    const readOutput = await env.DB.prepare(
+      'SELECT egress_bytes FROM retrieval_logs WHERE data_set_id = ?',
+    )
+      .bind(String(realDataSetId))
+      .all()
+    expect(readOutput.results).toStrictEqual([
+      expect.objectContaining({
+        egress_bytes: 4,
+      }),
+    ])
+  })
+
   it('requests payment if withCDN=false', async () => {
     const dataSetId = 'test-data-set-no-cdn'
     const pieceId = 'root-no-cdn'

piece-retriever/test/store.test.js

@@ -861,45 +861,4 @@ describe('updateDataSetStats', () => {
       initialCacheMissQuota - EGRESS_BYTES,
     )
   })
-
-  it('does not decrement quotas when isBotTraffic is true', async () => {
-    const DATA_SET_ID = 'test-data-set-bot-traffic'
-    const EGRESS_BYTES = 100
-    const initialCdnQuota = 500
-    const initialCacheMissQuota = 300
-
-    await withDataSet(env, {
-      dataSetId: DATA_SET_ID,
-      cdnEgressQuota: initialCdnQuota,
-      cacheMissEgressQuota: initialCacheMissQuota,
-    })
-
-    await updateDataSetStats(env, {
-      dataSetId: DATA_SET_ID,
-      egressBytes: EGRESS_BYTES,
-      cacheMiss: false,
-      enforceEgressQuota: true,
-      isBotTraffic: true,
-    })
-
-    const quotaResult = await env.DB.prepare(
-      'SELECT * FROM data_set_egress_quotas WHERE data_set_id = ?',
-    )
-      .bind(DATA_SET_ID)
-      .first()
-
-    expect(quotaResult).toStrictEqual({
-      data_set_id: DATA_SET_ID,
-      cdn_egress_quota: initialCdnQuota,
-      cache_miss_egress_quota: initialCacheMissQuota,
-    })
-
-    const dataSetResult = await env.DB.prepare(
-      'SELECT total_egress_bytes_used FROM data_sets WHERE id = ?',
-    )
-      .bind(DATA_SET_ID)
-      .first()
-
-    expect(dataSetResult.total_egress_bytes_used).toBe(EGRESS_BYTES)
-  })
 })