ci: fix Docker build step with tags

Author: Pouyanpi

.github/workflows/test-docker.yml

@@ -3,39 +3,54 @@ name: Test Docker Image
 # TODO: set docker auth
 on:
   workflow_dispatch:
-# on:
-#   push:
-#     tags:
-#       - "v*"
-
-env:
-  IMAGE: nemoguardrails
-  TEST_TAG: ${{ env.IMAGE }}:test
-  LATEST_TAG: ${{ env.IMAGE }}:latest
+  push:
+    #   tags:
+    #     - "v*"
+  pull_request:
 
 jobs:
   docker:
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        # os: [ubuntu-latest, macos-15, windows-latest]
+        os: [ubuntu-latest]
+    env:
+      IMAGE: nemoguardrails
+      GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
     steps:
       #
       # Checkout the code
       - name: Checkout
         uses: actions/checkout@v4
 
-      # Extract the tag version and set Docker tags
-      - name: Get git tag and set Docker tags
+      # Get runner architecture
+      - name: Get runner architecture
+        id: runner-arch
+        run: echo "arch=$(uname -m)" >> $GITHUB_OUTPUT
+
+      # Extract the PR SHA and set Docker tags
+      - name: Get PR SHA and set Docker tags
         run: |
-          VERSION=${GITHUB_REF#refs/tags/v}
-          echo "VERSION=$VERSION" >> $GITHUB_ENV
-          echo "VERSION_TAG=${{ env.IMAGE }}:$VERSION" >> $GITHUB_ENV
+          SHA=${GITHUB_SHA}
+          OS=${{ matrix.os }}
+          ARCH=${{ steps.runner-arch.outputs.arch }}
+          echo "SHA=$SHA" >> $GITHUB_ENV
+          echo "PR_TAG=${{ env.IMAGE }}:$SHA-$OS-$ARCH" >> $GITHUB_ENV
 
-      # Build the Docker image
-      - name: Build the Docker image
-        run: docker build . --file Dockerfile --tag ${{ env.TEST_TAG }} --tag ${{ env.VERSION_TAG }}
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      - name: Build image
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: Dockerfile
+          load: true
+          tags: ${{ env.PR_TAG }}
 
       # Start the container in detached mode
       - name: Start container
-        run: docker run -d --name test_container -p 8000:8000 ${{ env.TEST_TAG }}
+        run: docker run -d --name test_container -p 8000:8000 ${{ env.PR_TAG }}
 
       # Wait for the container to be ready
       - name: Wait for container to be ready