diff --git a/lms/djangoapps/edxnotes/decorators.py b/lms/djangoapps/edxnotes/decorators.py
index c641bd230c..247bb6b913 100644
--- a/lms/djangoapps/edxnotes/decorators.py
+++ b/lms/djangoapps/edxnotes/decorators.py
@@ -33,8 +33,8 @@ def edxnotes(cls):
         if not hasattr(runtime, 'modulestore'):
             return original_get_html(self, *args, **kwargs)
 
-        is_studio = getattr(self.system, "is_author_mode", False)
-        course = getattr(self, 'descriptor', self).runtime.modulestore.get_course(self.runtime.course_id)
+        is_studio = getattr(self.runtime, "is_author_mode", False)
+        course = getattr(self, 'descriptor', self).runtime.modulestore.get_course(self.scope_ids.usage_id.context_key)
 
         # Must be disabled when:
         # - in Studio
@@ -57,10 +57,10 @@ def edxnotes(cls):
                 ),
                 "params": {
                     # Use camelCase to name keys.
-                    "usageId": str(self.scope_ids.usage_id),
-                    "courseId": str(self.runtime.course_id),
+                    "usageId": self.scope_ids.usage_id,
+                    "courseId": course.id,
                     "token": get_edxnotes_id_token(user),
-                    "tokenUrl": get_token_url(self.runtime.course_id),
+                    "tokenUrl": get_token_url(course.id),
                     "endpoint": get_public_endpoint(),
                     "debug": settings.DEBUG,
                     "eventStringLimit": settings.TRACK_MAX_EVENT / 6,
diff --git a/lms/djangoapps/edxnotes/tests.py b/lms/djangoapps/edxnotes/tests.py
index a67ba7e0fa..b89b4a0afe 100644
--- a/lms/djangoapps/edxnotes/tests.py
+++ b/lms/djangoapps/edxnotes/tests.py
@@ -79,11 +79,10 @@ class TestProblem:
     The purpose of this class is to imitate any problem.
     """
     def __init__(self, course, user=None):
-        self.system = MagicMock(is_author_mode=False)
-        self.scope_ids = MagicMock(usage_id="test_usage_id")
+        self.scope_ids = MagicMock(usage_id=course.id.make_usage_key('test_problem', 'test_usage_id'))
         user = user or UserFactory()
         user_service = StubUserService(user)
-        self.runtime = MagicMock(course_id=course.id, service=lambda _a, _b: user_service)
+        self.runtime = MagicMock(service=lambda _a, _b: user_service, is_author_mode=False)
         self.descriptor = MagicMock()
         self.descriptor.runtime.modulestore.get_course.return_value = course
 
@@ -136,7 +135,7 @@ class EdxNotesDecoratorTest(ModuleStoreTestCase):
             "uid": "uid",
             "edxnotes_visibility": "true",
             "params": {
-                "usageId": "test_usage_id",
+                "usageId": problem.scope_ids.usage_id,
                 "courseId": course.id,
                 "token": "token",
                 "tokenUrl": "/tokenUrl",
@@ -167,7 +166,7 @@ class EdxNotesDecoratorTest(ModuleStoreTestCase):
         """
         Tests that get_html is not wrapped when problem is rendered in Studio.
         """
-        self.problem.system.is_author_mode = True
+        self.problem.runtime.is_author_mode = True
         assert 'original_get_html' == self.problem.get_html()
 
     def test_edxnotes_blockstore_runtime(self):