diff --git a/file_example/src/StreamWrapper/MockSessionTrait.php b/file_example/src/StreamWrapper/MockSessionTrait.php
new file mode 100644
index 0000000..4f62165
--- /dev/null
+++ b/file_example/src/StreamWrapper/MockSessionTrait.php
@@ -0,0 +1,101 @@
+<?php
+
+/**
+ * @file
+ * Contains \Drupal\file_example\StreamWrapper\MockSessionTrait.
+ */
+
+namespace Drupal\file_example\StreamWrapper;
+
+use Symfony\Component\HttpFoundation\RequestStack;
+use Symfony\Component\HttpFoundation\Request;
+use Symfony\Component\HttpFoundation\Session\SessionInterface;
+use Prophecy\Argument;
+use Drupal\file_example\StreamWrapper\SessionWrapper;
+
+trait MockSessionTrait {
+
+  /**
+   * @var array
+   *
+   * We'll use this to back our mock session.
+   */
+  protected $sessionStore;
+
+  /**
+   * @var RequestStack|ProphecyInterface
+   */
+  protected $requestStack;
+
+  /**
+   * Create a mock session object.
+   *
+   * @return ProphecyInterface
+   *   A test double, or mock, of a RequestStack object
+   *   that can be used to return a mock Session object.
+   */
+  protected function createSessionMock() {
+    $this->sessionStore = [];
+    $session = $this->prophesize(SessionInterface::class);
+    $test = $this;
+
+    $session
+      ->get('file_example', [])
+      ->will(function($args) use ($test) {
+        return $test->getSessionStore();
+      });
+
+    $session
+      ->set('file_example', Argument::any())
+      ->will(function($args) use ($test) {
+        $test->setSessionStore($args[1]);
+      });
+
+    $session
+      ->remove('file_example')
+      ->will(function($args) use ($test) {
+        $test->resetSessionStore();
+      });
+
+    $request = $this->prophesize(Request::class);
+    $request
+      ->getSession()
+      ->willReturn($session->reveal());
+
+    $request_stack = $this->prophesize(RequestStack::class);
+    $request_stack
+      ->getCurrentRequest()
+      ->willReturn($request->reveal());
+
+    return $this->requestStack = $request_stack->reveal();
+  }
+
+  /**
+   * Get a session wrapper.
+   */
+  public function getSessionWrapper() {
+    return new SessionWrapper($this->requestStack);
+  }
+
+  /**
+   * Helper for mocks.
+   */
+  public function getSessionStore() {
+    return $this->sessionStore;
+  }
+
+  /**
+   * Helper for our mocks.
+   */
+  public function setSessionStore($data) {
+    $this->sessionStore = $data;
+  }
+
+  /**
+   * Helper for our mocks.
+   */
+  public function resetSessionStore() {
+    $this->sessionStore = [];
+  }
+
+}
diff --git a/file_example/tests/src/Kernel/StreamWrapperTest.php b/file_example/tests/src/Kernel/StreamWrapperTest.php
index cf74de7..9bb9125 100644
--- a/file_example/tests/src/Kernel/StreamWrapperTest.php
+++ b/file_example/tests/src/Kernel/StreamWrapperTest.php
@@ -11,7 +11,7 @@ use Drupal\Component\FileCache\FileCacheFactory;
 use Drupal\Core\Site\Settings;
 use Drupal\KernelTests\KernelTestBase;
 use Drupal\Component\Utility\Html;
-use Drupal\Tests\file_example\MockSessionTrait;
+use Drupal\file_example\StreamWrapper\MockSessionTrait;
 use Drupal\Core\DependencyInjection\ContainerBuilder;
 
 /**
diff --git a/file_example/tests/src/MockSessionTrait.php b/file_example/tests/src/MockSessionTrait.php
deleted file mode 100644
index 1eeb361..0000000
--- a/file_example/tests/src/MockSessionTrait.php
+++ /dev/null
@@ -1,101 +0,0 @@
-<?php
-
-/**
- * @file
- * Contains \Drupal\Tests\file_example\MockSessionTrait.
- */
-
-namespace Drupal\Tests\file_example;
-
-use Symfony\Component\HttpFoundation\RequestStack;
-use Symfony\Component\HttpFoundation\Request;
-use Symfony\Component\HttpFoundation\Session\SessionInterface;
-use Prophecy\Argument;
-use Drupal\file_example\StreamWrapper\SessionWrapper;
-
-trait MockSessionTrait {
-
-  /**
-   * @var array
-   *
-   * We'll use this to back our mock session.
-   */
-  protected $sessionStore;
-
-  /**
-   * @var RequestStack|ProphecyInterface
-   */
-  protected $requestStack;
-
-  /**
-   * Create a mock session object.
-   *
-   * @return ProphecyInterface
-   *   A test double, or mock, of a RequestStack object
-   *   that can be used to return a mock Session object.
-   */
-  protected function createSessionMock() {
-    $this->sessionStore = [];
-    $session = $this->prophesize(SessionInterface::class);
-    $test = $this;
-
-    $session
-      ->get('file_example', [])
-      ->will(function($args) use ($test) {
-        return $test->getSessionStore();
-      });
-
-    $session
-      ->set('file_example', Argument::any())
-      ->will(function($args) use ($test) {
-        $test->setSessionStore($args[1]);
-      });
-
-    $session
-      ->remove('file_example')
-      ->will(function($args) use ($test) {
-        $test->resetSessionStore();
-      });
-
-    $request = $this->prophesize(Request::class);
-    $request
-      ->getSession()
-      ->willReturn($session->reveal());
-
-    $request_stack = $this->prophesize(RequestStack::class);
-    $request_stack
-      ->getCurrentRequest()
-      ->willReturn($request->reveal());
-
-    return $this->requestStack = $request_stack->reveal();
-  }
-
-  /**
-   * Get a session wrapper.
-   */
-  public function getSessionWrapper() {
-    return new SessionWrapper($this->requestStack);
-  }
-
-  /**
-   * Helper for mocks.
-   */
-  public function getSessionStore() {
-    return $this->sessionStore;
-  }
-
-  /**
-   * Helper for our mocks.
-   */
-  public function setSessionStore($data) {
-    $this->sessionStore = $data;
-  }
-
-  /**
-   * Helper for our mocks.
-   */
-  public function resetSessionStore() {
-    $this->sessionStore = [];
-  }
-
-}
diff --git a/file_example/tests/src/Unit/SessionWrapperTest.php b/file_example/tests/src/Unit/SessionWrapperTest.php
index 3aa9537..0f85978 100644
--- a/file_example/tests/src/Unit/SessionWrapperTest.php
+++ b/file_example/tests/src/Unit/SessionWrapperTest.php
@@ -9,7 +9,8 @@ namespace Drupal\Tests\file_example\Unit;
 
 use Drupal\Tests\UnitTestCase;
 use Drupal\file_example\StreamWrapper\SessionWrapper;
-use Drupal\Tests\file_example\MockSessionTrait;
+use Drupal\file_example\StreamWrapper\MockSessionTrait;
+
 /**
  * PHPUnit test for the SessionWrapper session manipulation class.
  *