diff --git a/examples/java/src/test/java/dev/selenium/browsers/FirefoxTest.java b/examples/java/src/test/java/dev/selenium/browsers/FirefoxTest.java
index bc748391e2a..70ea5c18240 100644
--- a/examples/java/src/test/java/dev/selenium/browsers/FirefoxTest.java
+++ b/examples/java/src/test/java/dev/selenium/browsers/FirefoxTest.java
@@ -15,13 +15,13 @@
 import org.openqa.selenium.By;
 import org.openqa.selenium.OutputType;
 import org.openqa.selenium.WebElement;
-import org.openqa.selenium.firefox.FirefoxDriver;
-import org.openqa.selenium.firefox.FirefoxDriverLogLevel;
-import org.openqa.selenium.firefox.FirefoxDriverService;
-import org.openqa.selenium.firefox.FirefoxOptions;
-import org.openqa.selenium.firefox.GeckoDriverService;
+import org.openqa.selenium.firefox.*;
 import org.openqa.selenium.remote.service.DriverFinder;
 
+
+
+
+
 public class FirefoxTest extends BaseTest {
   private FirefoxDriver driver;
 
@@ -189,4 +189,20 @@ public void fullPageScreenshot() throws Exception {
 
     driver.quit();
   }
+
+  @Test
+  public void setContext() {
+    driver = startFirefoxDriver();
+
+    ((HasContext) driver).setContext(FirefoxCommandContext.CHROME);
+    driver.executeScript("console.log('Inside Chrome context');");
+
+    // Verify the context is back to "content"
+    Assertions.assertEquals(
+            FirefoxCommandContext.CHROME, ((HasContext) driver).getContext(),
+            "The context should be 'chrome'"
+    );
+
+    driver.quit();
+  }
 }
diff --git a/website_and_docs/content/documentation/webdriver/browsers/firefox.en.md b/website_and_docs/content/documentation/webdriver/browsers/firefox.en.md
index 2dd14391956..f2eac7c6b54 100644
--- a/website_and_docs/content/documentation/webdriver/browsers/firefox.en.md
+++ b/website_and_docs/content/documentation/webdriver/browsers/firefox.en.md
@@ -441,7 +441,7 @@ please refer to the
 
 {{< tabpane text=true >}}
 {{< tab header="Java" >}}
-{{< badge-code >}}
+{{< gh-codeblock path="/examples/java/src/test/java/dev/selenium/browsers/FirefoxTest.java#L197-L198" >}}
 {{< /tab >}}
 {{< tab header="Python" >}}
 {{< gh-codeblock path="/examples/python/tests/browsers/test_firefox.py#L149-L150" >}}
diff --git a/website_and_docs/content/documentation/webdriver/browsers/firefox.ja.md b/website_and_docs/content/documentation/webdriver/browsers/firefox.ja.md
index d489c0be1af..126638d5acb 100644
--- a/website_and_docs/content/documentation/webdriver/browsers/firefox.ja.md
+++ b/website_and_docs/content/documentation/webdriver/browsers/firefox.ja.md
@@ -437,7 +437,7 @@ IDはアドオンインストール時の戻り値から取得できます。
 
 {{< tabpane text=true >}}
 {{< tab header="Java" >}}
-{{< badge-code >}}
+{{< gh-codeblock path="/examples/java/src/test/java/dev/selenium/browsers/FirefoxTest.java#L197-L198" >}}
 {{< /tab >}}
 {{< tab header="Python" >}}
 {{< gh-codeblock path="/examples/python/tests/browsers/test_firefox.py#L149-L150" >}}
diff --git a/website_and_docs/content/documentation/webdriver/browsers/firefox.pt-br.md b/website_and_docs/content/documentation/webdriver/browsers/firefox.pt-br.md
index 01cc173f556..f36890f487d 100644
--- a/website_and_docs/content/documentation/webdriver/browsers/firefox.pt-br.md
+++ b/website_and_docs/content/documentation/webdriver/browsers/firefox.pt-br.md
@@ -446,7 +446,7 @@ please refer to the
 
 {{< tabpane text=true >}}
 {{< tab header="Java" >}}
-{{< badge-code >}}
+{{< gh-codeblock path="/examples/java/src/test/java/dev/selenium/browsers/FirefoxTest.java#L197-L198" >}}
 {{< /tab >}}
 {{< tab header="Python" >}}
 {{< gh-codeblock path="/examples/python/tests/browsers/test_firefox.py#L149-L150" >}}
diff --git a/website_and_docs/content/documentation/webdriver/browsers/firefox.zh-cn.md b/website_and_docs/content/documentation/webdriver/browsers/firefox.zh-cn.md
index a30d1c9e751..b56e5f62b74 100644
--- a/website_and_docs/content/documentation/webdriver/browsers/firefox.zh-cn.md
+++ b/website_and_docs/content/documentation/webdriver/browsers/firefox.zh-cn.md
@@ -443,7 +443,7 @@ please refer to the
 
 {{< tabpane text=true >}}
 {{< tab header="Java" >}}
-{{< badge-code >}}
+{{< gh-codeblock path="/examples/java/src/test/java/dev/selenium/browsers/FirefoxTest.java#L197-L198" >}}
 {{< /tab >}}
 {{< tab header="Python" >}}
 {{< gh-codeblock path="/examples/python/tests/browsers/test_firefox.py#L149-L150" >}}