diff --git a/__tests__/main.test.js b/__tests__/main.test.js
index 321a46d..31d1c85 100644
--- a/__tests__/main.test.js
+++ b/__tests__/main.test.js
@@ -1,26 +1,24 @@
-/**
- * Unit tests for the action's main functionality, src/main.js
- */
 const core = require('@actions/core')
-const { context } = require('@actions/github')
+const github = require('@actions/github')
+const { Octokit } = require('@octokit/core')
 const main = require('../src/main')
 
-// Mock the GitHub Actions core library
-const getInputMock = jest.spyOn(core, 'getInput').mockImplementation()
-const setFailedMock = jest.spyOn(core, 'setFailed').mockImplementation()
-const setOutputMock = jest.spyOn(core, 'setOutput').mockImplementation()
+jest.mock('@actions/core')
+jest.mock('@actions/github', () => ({
+  context: {
+    ref: 'refs/heads/test-branch'
+  }
+}))
+jest.mock('@octokit/core')
 
-// Mock the action's main function
-const runMock = jest.spyOn(main, 'run')
+describe('run function', () => {
+  let getInputMock
+  let setOutputMock
+  let setFailedMock
+  let octokitMock
 
-describe('action', () => {
   beforeEach(() => {
-    jest.clearAllMocks()
-  })
-
-  it('sets the correct outputs', async () => {
-    // Set the action's inputs as return values from core.getInput()
-    getInputMock.mockImplementation(name => {
+    getInputMock = core.getInput.mockImplementation(name => {
       switch (name) {
         case 'token':
           return 'test-token'
@@ -33,32 +31,68 @@ describe('action', () => {
       }
     })
 
-    // Mock github.context.ref to return a test branch name
-    context.ref = 'refs/heads/test-branch'
+    setOutputMock = core.setOutput.mockImplementation()
+    setFailedMock = core.setFailed.mockImplementation()
 
-    await main.run()
-    expect(runMock).toHaveReturned()
+    octokitMock = {
+      request: jest.fn()
+    }
 
-    // Verify that all of the core library functions were called correctly
-    // expect(setOutputMock).toHaveBeenNthCalledWith(1, 'time', expect.any(String))
+    Octokit.mockImplementation(() => octokitMock)
+
+    github.context.ref = 'refs/heads/test-branch'
   })
 
-  it('sets a failed status', async () => {
-    // Set the action's inputs as return values from core.getInput()
-    getInputMock.mockImplementation(name => {
-      switch (name) {
-        case 'token':
-          return 'test-token'
-        case 'owner':
-          return 'test-owner'
-        case 'repo':
-          return 'test-repo'
-        default:
-          return ''
+  afterEach(() => {
+    jest.clearAllMocks()
+  })
+
+  it('should set the correct outputs', async () => {
+    octokitMock.request.mockResolvedValueOnce({
+      data: [
+        {
+          id: 1,
+          ref: 'fix/switch-header-and-footer'
+        }
+      ]
+    })
+    octokitMock.request.mockResolvedValueOnce({
+      status: 200
+    })
+
+    await main.run()
+
+    expect(getInputMock).toHaveBeenCalledWith('token', { required: true })
+    expect(getInputMock).toHaveBeenCalledWith('owner', { required: true })
+    expect(getInputMock).toHaveBeenCalledWith('repo', { required: true })
+    expect(octokitMock.request).toHaveBeenCalledWith(
+      'GET /repos/{owner}/{repo}/deployments',
+      {
+        owner: 'test-owner',
+        repo: 'test-repo'
+      }
+    )
+    expect(octokitMock.request).toHaveBeenCalledWith(
+      'POST /repos/{owner}/{repo}/deployments/{id}/statuses',
+      {
+        owner: 'test-owner',
+        repo: 'test-repo',
+        id: 1,
+        data: { state: 'inactive' }
       }
+    )
+    expect(setOutputMock).toHaveBeenCalledWith('time', expect.any(String))
+  })
+
+  it('should fail when no deployment is found for the branch', async () => {
+    octokitMock.request.mockResolvedValueOnce({
+      data: []
     })
 
     await main.run()
-    expect(runMock).toHaveReturned()
+
+    expect(setFailedMock).toHaveBeenCalledWith(
+      'No deployment found for branch: fix/switch-header-and-footer'
+    )
   })
 })
diff --git a/src/main.js b/src/main.js
index 9a41171..908ffeb 100644
--- a/src/main.js
+++ b/src/main.js
@@ -10,8 +10,7 @@ async function listDeployments(octokit, owner, repo) {
       repo
     }
   )
-
-  return response.data
+  return response.data || []
 }
 
 async function changeStatusDeployment(octokit, owner, repo, id) {
@@ -32,44 +31,32 @@ async function changeStatusDeployment(octokit, owner, repo, id) {
   return response.status === 200
 }
 
-/**
- * The main function for the action.
- * @returns {Promise<void>} Resolves when the action is complete.
- */
 async function run() {
   try {
     const token = core.getInput('token', { required: true })
     const owner = core.getInput('owner', { required: true })
     const repo = core.getInput('repo', { required: true })
 
-    // Get the branch that triggered the action
     const ref = github.context.ref
-    // const branch = ref.replace('refs/heads/', '')
-
     const branch = 'fix/switch-header-and-footer'
 
-    // Log the branch name to the console
     console.log(`Branch: ${branch}`)
 
-    // Request for github
     const octokit = new Octokit({ auth: token })
 
-    // List deployments for the repository
-    const deployments = listDeployments(octokit, owner, repo)
+    const deployments = await listDeployments(octokit, owner, repo)
     const currentBranchDeployment = deployments.find(deployment => {
       console.log('deployment', deployment.ref)
       console.log('branch', branch)
 
-      if (deployment.ref === branch) {
-        return deployment
-      }
+      return deployment.ref === branch
     })
 
     if (!currentBranchDeployment) {
       throw new Error(`No deployment found for branch: ${branch}`)
     }
 
-    const result = changeStatusDeployment(
+    const result = await changeStatusDeployment(
       octokit,
       owner,
       repo,
@@ -78,14 +65,14 @@ async function run() {
 
     console.log('Result change status', result)
 
-    // Set outputs for other workflow steps to use
     core.setOutput('time', new Date().toTimeString())
   } catch (error) {
-    // Fail the workflow run if an error occurs
     core.setFailed(error.message)
   }
 }
 
 module.exports = {
-  run
+  run,
+  listDeployments,
+  changeStatusDeployment
 }