chore: Add tests for ProviderDatabase and evictOrphanedEntityProviders

Signed-off-by: Jack Palmer <jackpalmer@spotify.com>

plugins/catalog-backend/config.d.ts

@@ -218,10 +218,7 @@ export interface Config {
     useUrlReadersSearch?: boolean;
 
     /**
-     * Evicts entities from the catalog that are no longer referenced by any
-     * added entity providers.
-     *
-     * Defaults to false.
+     * Evicts entities from the catalog that are no longer referenced by entity providers added to the catalog.
      */
     evictOrphanedEntityProviders?: boolean;
   };

plugins/catalog-backend/src/database/DefaultProviderDatabase.test.ts

@@ -59,21 +59,21 @@ describe('DefaultProviderDatabase', () => {
     await db<DbRefreshStateRow>('refresh_state').insert(ref);
   };
 
-  describe('replaceUnprocessedEntities', () => {
-    const createLocations = async (db: Knex, entityRefs: string[]) => {
-      for (const ref of entityRefs) {
-        await insertRefreshStateRow(db, {
-          entity_id: uuid.v4(),
-          entity_ref: ref,
-          unprocessed_entity: '{}',
-          processed_entity: '{}',
-          errors: '[]',
-          next_update_at: '2021-04-01 13:37:00',
-          last_discovery_at: '2021-04-01 13:37:00',
-        });
-      }
-    };
+  const createLocations = async (db: Knex, entityRefs: string[]) => {
+    for (const ref of entityRefs) {
+      await insertRefreshStateRow(db, {
+        entity_id: uuid.v4(),
+        entity_ref: ref,
+        unprocessed_entity: '{}',
+        processed_entity: '{}',
+        errors: '[]',
+        next_update_at: '2021-04-01 13:37:00',
+        last_discovery_at: '2021-04-01 13:37:00',
+      });
+    }
+  };
 
+  describe('replaceUnprocessedEntities', () => {
     it.each(databases.eachSupportedId())(
       'replaces all existing state correctly for simple dependency chains, %p',
       async databaseId => {
@@ -988,4 +988,86 @@ describe('DefaultProviderDatabase', () => {
       },
     );
   });
+
+  describe('listReferenceSourceKeys', () => {
+    it.each(databases.eachSupportedId())(
+      'returns the source_keys from "refresh_state_references", %p',
+      async databaseId => {
+        const { knex, db } = await createDatabase(databaseId);
+
+        await createLocations(knex, [
+          'location:default/root',
+          'location:default/root-1',
+        ]);
+
+        await insertRefRow(knex, {
+          source_key: 'foo',
+          target_entity_ref: 'location:default/root',
+        });
+        await insertRefRow(knex, {
+          source_key: 'bar',
+          target_entity_ref: 'location:default/root-1',
+        });
+
+        const res = await db.transaction(async tx =>
+          db.listReferenceSourceKeys(tx),
+        );
+
+        expect(res).toEqual(['bar', 'foo']);
+      },
+    );
+
+    it.each(databases.eachSupportedId())(
+      'returns only unique source_keys", %p',
+      async databaseId => {
+        const { knex, db } = await createDatabase(databaseId);
+
+        await createLocations(knex, [
+          'location:default/root',
+          'location:default/root-1',
+        ]);
+
+        await insertRefRow(knex, {
+          source_key: 'foo',
+          target_entity_ref: 'location:default/root',
+        });
+        await insertRefRow(knex, {
+          source_key: 'foo',
+          target_entity_ref: 'location:default/root-1',
+        });
+
+        const res = await db.transaction(async tx =>
+          db.listReferenceSourceKeys(tx),
+        );
+
+        expect(res).toEqual(['foo']);
+      },
+    );
+
+    it.each(databases.eachSupportedId())(
+      'does not return null source_keys", %p',
+      async databaseId => {
+        const { knex, db } = await createDatabase(databaseId);
+
+        await createLocations(knex, [
+          'location:default/root',
+          'location:default/root-1',
+        ]);
+
+        await insertRefRow(knex, {
+          source_key: 'foo',
+          target_entity_ref: 'location:default/root',
+        });
+        await insertRefRow(knex, {
+          target_entity_ref: 'location:default/root-1',
+        });
+
+        const res = await db.transaction(async tx =>
+          db.listReferenceSourceKeys(tx),
+        );
+
+        expect(res).toEqual(['foo']);
+      },
+    );
+  });
 });

plugins/catalog-backend/src/database/DefaultProviderDatabase.ts

@@ -199,7 +199,7 @@ export class DefaultProviderDatabase implements ProviderDatabase {
     }
   }
 
-  async listEntityProviderNames(txOpaque: Transaction): Promise<string[]> {
+  async listReferenceSourceKeys(txOpaque: Transaction): Promise<string[]> {
     const tx = txOpaque as Knex | Knex.Transaction;
 
     const rows = await tx<DbRefreshStateReferencesRow>(

plugins/catalog-backend/src/database/types.ts

@@ -177,7 +177,7 @@ export interface ProviderDatabase {
   /**
    * List the names of all the entity providers that have references in the provider database.
    */
-  listEntityProviderNames(txOpaque: Transaction): Promise<string[]>;
+  listReferenceSourceKeys(txOpaque: Transaction): Promise<string[]>;
 }
 
 // TODO(Rugvip): This is only partial for now

plugins/catalog-backend/src/processing/evictOrphanedEntityProviders.test.ts

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2025 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import { EntityProvider } from '@backstage/plugin-catalog-node';
+import { mockServices } from '@backstage/backend-test-utils';
+import { DefaultProviderDatabase } from '../database/DefaultProviderDatabase';
+import { evictOrphanedEntityProviders } from './evictOrphanedEntityProviders';
+
+describe('evictOrphanedEntityProviders', () => {
+  const db = {
+    transaction: jest.fn().mockImplementation(cb => cb((() => {}) as any)),
+    replaceUnprocessedEntities: jest.fn(),
+    listReferenceSourceKeys: jest.fn(),
+  } as unknown as jest.Mocked<DefaultProviderDatabase>;
+
+  const providers = [
+    { getProviderName: () => 'provider1' },
+    { getProviderName: () => 'provider2' },
+  ] as unknown as EntityProvider[];
+  const logger = mockServices.logger.mock();
+
+  it('replaces unprocessed entities for orphaned providers with empty items', async () => {
+    db.listReferenceSourceKeys.mockResolvedValue(['foo', 'bar']);
+
+    await evictOrphanedEntityProviders({ db, providers, logger });
+
+    expect(db.replaceUnprocessedEntities).toHaveBeenCalledTimes(2);
+    expect(db.replaceUnprocessedEntities).toHaveBeenNthCalledWith(
+      1,
+      expect.anything(),
+      {
+        sourceKey: 'foo',
+        type: 'full',
+        items: [],
+      },
+    );
+    expect(db.replaceUnprocessedEntities).toHaveBeenNthCalledWith(
+      2,
+      expect.anything(),
+      {
+        sourceKey: 'bar',
+        type: 'full',
+        items: [],
+      },
+    );
+  });
+
+  it('does not replace unprocessed entities for providers that are not orphaned', async () => {
+    db.listReferenceSourceKeys.mockResolvedValue(['foo', 'provider1']);
+
+    await evictOrphanedEntityProviders({ db, providers, logger });
+
+    expect(db.replaceUnprocessedEntities).not.toHaveBeenCalledWith(
+      expect.anything(),
+      {
+        sourceKey: 'provider1',
+        type: 'full',
+        items: [],
+      },
+    );
+  });
+});

plugins/catalog-backend/src/processing/evictOrphanedEntityProviders.ts

@@ -15,15 +15,18 @@
  */
 
 import { EntityProvider } from '@backstage/plugin-catalog-node';
-import { ProviderDatabase } from '../database/types';
 import { LoggerService } from '@backstage/backend-plugin-api';
+import { ProviderDatabase } from '../database/types';
 
-async function getOrphanedEntityProviderNames(
-  db: ProviderDatabase,
-  providers: EntityProvider[],
-): Promise<string[]> {
+async function getOrphanedEntityProviderNames({
+  db,
+  providers,
+}: {
+  db: ProviderDatabase;
+  providers: EntityProvider[];
+}): Promise<string[]> {
   const dbProviderNames = await db.transaction(async tx =>
-    db.listEntityProviderNames(tx),
+    db.listReferenceSourceKeys(tx),
   );
 
   const providerNames = providers.map(p => p.getProviderName());
@@ -33,11 +36,15 @@ async function getOrphanedEntityProviderNames(
   );
 }
 
-async function removeEntitiesForProvider(
-  db: ProviderDatabase,
-  providerName: string,
-  logger: LoggerService,
-) {
+async function removeEntitiesForProvider({
+  db,
+  providerName,
+  logger,
+}: {
+  db: ProviderDatabase;
+  providerName: string;
+  logger: LoggerService;
+}) {
   try {
     await db.transaction(async tx => {
       await db.replaceUnprocessedEntities(tx, {
@@ -55,15 +62,16 @@ async function removeEntitiesForProvider(
   }
 }
 
-export async function evictOrphanedEntityProviders(
-  db: ProviderDatabase,
-  providers: EntityProvider[],
-  logger: LoggerService,
-) {
-  for (const providerName of await getOrphanedEntityProviderNames(
-    db,
-    providers,
-  )) {
-    await removeEntitiesForProvider(db, providerName, logger);
+export async function evictOrphanedEntityProviders(options: {
+  db: ProviderDatabase;
+  providers: EntityProvider[];
+  logger: LoggerService;
+}) {
+  for (const providerName of await getOrphanedEntityProviderNames(options)) {
+    await removeEntitiesForProvider({
+      db: options.db,
+      providerName,
+      logger: options.logger,
+    });
   }
 }

plugins/catalog-backend/src/service/CatalogBuilder.ts

@@ -647,11 +647,11 @@ export class CatalogBuilder {
     await connectEntityProviders(providerDatabase, entityProviders);
 
     if (config.getOptionalBoolean('catalog.evictOrphanedEntityProviders')) {
-      await evictOrphanedEntityProviders(
-        providerDatabase,
-        entityProviders,
+      await evictOrphanedEntityProviders({
+        db: providerDatabase,
+        providers: entityProviders,
         logger,
-      );
+      });
     }
 
     return {