diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 6d9df3844..1ae55a779 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -346,7 +346,7 @@ jobs:
     uses: ./.github/workflows/publish-package.yml
     with:
       publish: false
-      repository: testpypi
+      environment: test
       # don't have access to env context here for some reason
       ref: ${{ github.event_name == 'workflow_dispatch' && inputs.ref || null }}
       # can't use env context here so need to duplicate expression, but these are true boolean values so don't need extra string logic
diff --git a/.github/workflows/publish-documentation.yml b/.github/workflows/publish-documentation.yml
index 499764cd6..0f9731433 100644
--- a/.github/workflows/publish-documentation.yml
+++ b/.github/workflows/publish-documentation.yml
@@ -91,11 +91,12 @@ jobs:
         
   publish-docs:
     name: Publish documentation
-    runs-on: ubuntu-latest
     needs: create_docs
     permissions:
       id-token: write # needed to publish to Azure
+    environment: ${{ inputs.environment }}
     if: ${{ inputs.publish }}
+    runs-on: ubuntu-latest
     steps:
       - name: Download docs artifact
         uses: actions/download-artifact@v3
diff --git a/.github/workflows/publish-package.yml b/.github/workflows/publish-package.yml
index 7cd4322c2..d9dfbbdca 100644
--- a/.github/workflows/publish-package.yml
+++ b/.github/workflows/publish-package.yml
@@ -8,12 +8,12 @@ on:
         required: false
         default: true
         type: boolean
-      repository:
+      environment:
         description: 'Whether to publish to production PyPI or test PyPI'
         required: false
-        default: pypi
+        default: prod
         type: choice
-        options: [pypi, testpypi]
+        options: [prod, test]
       ref:
         description: 'The git ref to build the package for'
         required: false
@@ -34,10 +34,10 @@ on:
         default: true
         type: boolean
       # choice type only supported for workflow_dispatch, not workflow_call
-      repository:
+      environment:
         description: 'Whether to publish to production PyPI or test PyPI'
         required: false
-        default: pypi
+        default: prod
         type: string
       ref:
         description: 'The git ref to build the package for'
@@ -119,6 +119,7 @@ jobs:
     needs: [merge]
     permissions:
       id-token: write
+    environment: ${{ inputs.environment }}
     if: ${{ inputs.publish }}
     runs-on: ubuntu-latest
     steps:
@@ -130,4 +131,4 @@ jobs:
       - name: Upload wheels and sdist to package index
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
-          repository-url: ${{ inputs.repository == 'testpypi' && 'https://test.pypi.org/legacy/' || '' }}
+          repository-url: ${{ inputs.environment == 'test' && 'https://test.pypi.org/legacy/' || '' }}