diff --git a/packages/api/src/resolvers/function_resolvers.ts b/packages/api/src/resolvers/function_resolvers.ts
index 00b63e8f9..7b5b1c6e9 100644
--- a/packages/api/src/resolvers/function_resolvers.ts
+++ b/packages/api/src/resolvers/function_resolvers.ts
@@ -583,5 +583,4 @@ export const functionResolvers = {
   ...resultResolveTypeResolver('Webhook'),
   ...resultResolveTypeResolver('ApiKeys'),
   ...resultResolveTypeResolver('RevokeApiKey'),
-  ...resultResolveTypeResolver('SetLabelsForHighlight'),
 }
diff --git a/packages/api/test/resolvers/labels.test.ts b/packages/api/test/resolvers/labels.test.ts
index 6d7814fc7..0701cdc0b 100644
--- a/packages/api/test/resolvers/labels.test.ts
+++ b/packages/api/test/resolvers/labels.test.ts
@@ -9,11 +9,12 @@ import { Label } from '../../src/entity/label'
 import { expect } from 'chai'
 import 'mocha'
 import { User } from '../../src/entity/user'
-import { Page, PageContext } from '../../src/elastic/types'
+import { Highlight, Page, PageContext } from '../../src/elastic/types'
 import { getRepository } from '../../src/entity/utils'
 import { getPageById } from '../../src/elastic/pages'
 import { addLabelInPage } from '../../src/elastic/labels'
 import { createPubSubClient } from '../../src/datalayer/pubsub'
+import { addHighlightToPage } from '../../src/elastic/highlights'
 
 describe('Labels API', () => {
   const username = 'fakeUser'
@@ -440,4 +441,96 @@ describe('Labels API', () => {
       })
     })
   })
+
+  describe('Set labels for highlight', () => {
+    let query: string
+    let highlightId: string
+    let labelIds: string[] = []
+
+    beforeEach(() => {
+      query = `
+        mutation {
+          setLabelsForHighlight(
+            input: {
+              highlightId: "${highlightId}",
+              labelIds: [
+                "${labelIds[0]}",
+                "${labelIds[1]}"
+              ]
+            }
+          ) {
+            ... on SetLabelsSuccess {
+              labels {
+                id
+                name
+              }
+            }
+            ... on SetLabelsError {
+              errorCodes
+            }
+          }
+        }
+      `
+    })
+
+    context('when labels exists', () => {
+      before(async () => {
+        highlightId = 'highlight-id'
+        const highlight: Highlight = {
+          createdAt: new Date(),
+          id: highlightId,
+          patch: 'test patch',
+          quote: 'test quote',
+          shortId: 'test shortId',
+          userId: user.id,
+        }
+        await addHighlightToPage(page.id, highlight, ctx)
+        labelIds = [labels[0].id, labels[1].id]
+      })
+
+      it('should set labels for highlight', async () => {
+        const res = await graphqlRequest(query, authToken).expect(200)
+        expect(
+          res.body.data.setLabelsForHighlight.labels.map((l: any) => l.id)
+        ).to.eql(labelIds)
+      })
+    })
+
+    context('when labels not exist', () => {
+      before(async () => {
+        highlightId = 'highlight-id-2'
+        const highlight: Highlight = {
+          createdAt: new Date(),
+          id: highlightId,
+          patch: 'test patch',
+          quote: 'test quote',
+          shortId: 'test shortId',
+          userId: user.id,
+        }
+        await addHighlightToPage(page.id, highlight, ctx)
+        labelIds = [generateFakeUuid(), generateFakeUuid()]
+      })
+
+      it('should return error code NOT_FOUND', async () => {
+        const res = await graphqlRequest(query, authToken).expect(200)
+        expect(res.body.data.setLabelsForHighlight.errorCodes).to.eql([
+          'NOT_FOUND',
+        ])
+      })
+    })
+
+    context('when highlight not exist', () => {
+      before(() => {
+        highlightId = generateFakeUuid()
+        labelIds = [labels[0].id, labels[1].id]
+      })
+
+      it('should return error code NOT_FOUND', async () => {
+        const res = await graphqlRequest(query, authToken).expect(200)
+        expect(res.body.data.setLabelsForHighlight.errorCodes).to.eql([
+          'NOT_FOUND',
+        ])
+      })
+    })
+  })
 })