diff --git a/frontend/src/app/models/found-document.spec.ts b/frontend/src/app/models/found-document.spec.ts
index 7f8935b3a..161ee64f2 100644
--- a/frontend/src/app/models/found-document.spec.ts
+++ b/frontend/src/app/models/found-document.spec.ts
@@ -67,9 +67,9 @@ describe('FoundDocument', () => {
         const doc = makeDocument({ great_field: 'test' });
         expect(doc.tags$.value).toEqual(mockTags);
         const tag = _.first(mockTags);
-        doc.removeTag(tag.id);
+        doc.removeTag(tag);
         expect(doc.tags$.value.length).toBe(1);
-        doc.addTag(tag.id);
+        doc.addTag(tag);
         expect(doc.tags$.value.length).toBe(2);
     });
 });
diff --git a/frontend/src/mock-data/tag.ts b/frontend/src/mock-data/tag.ts
index 2195e17fc..7d08e7a13 100644
--- a/frontend/src/mock-data/tag.ts
+++ b/frontend/src/mock-data/tag.ts
@@ -29,8 +29,8 @@ export class TagServiceMock {
         }).pipe(tap(this.fetch.bind(this)));
     }
 
-    setDocumentTags(document: FoundDocument, tagIds: number[]): Observable<Tag[]> {
-        const tags = mockTags.filter(tag => tagIds.includes(tag.id));
+    setDocumentTags(document: FoundDocument, tagIds: Tag[]): Observable<Tag[]> {
+        const tags = mockTags.filter(tag => tagIds.includes(tag));
         return of(tags);
     };