diff --git a/.github/actions/test/action.yml b/.github/actions/test/action.yml
new file mode 100644
index 0000000..17cfc7f
--- /dev/null
+++ b/.github/actions/test/action.yml
@@ -0,0 +1,51 @@
+name: 'rai-sdk-go test'
+description: 'rai-sdk-go test action'
+
+inputs:
+  client_id:
+    required: true
+    description: 'Client ID for oAuth'
+
+  client_secret:
+    required: true
+    description: 'Client secret for oAuth'
+
+  client_credentials_url:
+    required: true
+    description: 'Client credentials url for fetching the oAuth token'
+
+  rai_host:
+    required: false
+    description: 'RAI host'
+    default: 'azure.relationalai.com'
+
+  custom_headers:
+    required: false
+    description: 'Optional http headers'
+    default: '{}'
+
+runs:
+  using: 'composite'
+  steps:
+    - uses: actions/checkout@v3
+      with:
+        repository: RelationalAI/rai-sdk-go
+
+    - name: Set up Go
+      uses: actions/setup-go@v2
+      with:
+        go-version: 1.18
+
+    - name: Build
+      run: go build -v ./rai
+      shell: bash
+
+    - name: Test
+      env:
+        CLIENT_ID: ${{ inputs.client_id }}
+        CLIENT_SECRET: ${{ inputs.client_secret }}
+        CLIENT_CREDENTIALS_URL: ${{ inputs.client_credentials_url }}
+        HOST: ${{ inputs.rai_host }}
+        CUSTOM_HEADERS: ${{ inputs.custom_headers }}
+      run: go test -v -timeout 30m ./rai
+      shell: bash
diff --git a/.github/workflows/go-build.yaml b/.github/workflows/go-build.yaml
index 5122784..b4b6e12 100644
--- a/.github/workflows/go-build.yaml
+++ b/.github/workflows/go-build.yaml
@@ -9,19 +9,11 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-
-    - name: Set up Go
-      uses: actions/setup-go@v2
-      with:
-        go-version: 1.18
-
-    - name: Build
-      run: go build -v ./rai
+    - uses: actions/checkout@v3
 
     - name: Test
-      env:
-          CLIENT_ID: ${{ secrets.CLIENT_ID }}
-          CLIENT_SECRET: ${{ secrets.CLIENT_SECRET }}
-          CLIENT_CREDENTIALS_URL: ${{ secrets.CLIENT_CREDENTIALS_URL }}
-      run: go test -v -timeout 30m ./rai
+      uses: ./.github/actions/test
+      with:
+        client_id: ${{ secrets.CLIENT_ID }}
+        client_secret: ${{ secrets.CLIENT_SECRET }}
+        client_credentials_url: ${{ secrets.CLIENT_CREDENTIALS_URL }}
diff --git a/rai/main_test.go b/rai/main_test.go
index 1044c33..6fd97d8 100644
--- a/rai/main_test.go
+++ b/rai/main_test.go
@@ -100,10 +100,14 @@ func newTestClient() (*Client, error) {
 	clientId := os.Getenv("CLIENT_ID")
 	clientSecret := os.Getenv("CLIENT_SECRET")
 	clientCredentialsUrl := os.Getenv("CLIENT_CREDENTIALS_URL")
+	raiHost := os.Getenv("HOST")
+	if raiHost == "" {
+		raiHost = "azure.relationalai.com"
+	}
 
 	placeHolderConfig := `
 		[default]
-		host=azure.relationalai.com
+		host=%s
 		region=us-east
 		port=443
 		scheme=https
@@ -111,7 +115,7 @@ func newTestClient() (*Client, error) {
 		client_secret=%s
 		client_credentials_url=%s
 	`
-	configSrc := fmt.Sprintf(placeHolderConfig, clientId, clientSecret, clientCredentialsUrl)
+	configSrc := fmt.Sprintf(placeHolderConfig, raiHost, clientId, clientSecret, clientCredentialsUrl)
 	LoadConfigString(configSrc, "default", &cfg)
 	opts := ClientOptions{Config: cfg}
 	testClient := NewClient(context.Background(), &opts)