From 15829b90015f17220881f8dab54e3961a341f11c Mon Sep 17 00:00:00 2001
From: Michael Wheeler <mpwheel@umich.edu>
Date: Mon, 25 Mar 2024 12:20:52 -0400
Subject: [PATCH] Fix membership retrieval (#95)

The aud claim may be a string or an array.
---
 lti_tool/models.py   |  3 ++-
 tests/test_models.py | 22 ++++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/lti_tool/models.py b/lti_tool/models.py
index 3043ba0..8f3e8c8 100644
--- a/lti_tool/models.py
+++ b/lti_tool/models.py
@@ -627,7 +627,8 @@ def roles_claim(self):
     def membership(self) -> LtiMembership:
         """The LTI membership associated with the launch."""
         issuer = self.get_claim("iss")
-        client_id = self.get_claim("aud")
+        aud = self.get_claim("aud")
+        client_id = aud[0] if isinstance(aud, list) else aud
         sub = self.get_claim("sub")
         context_id = self.context_claim["id"] if self.context_claim is not None else ""
         return LtiMembership.objects.get(
diff --git a/tests/test_models.py b/tests/test_models.py
index b11becb..442f458 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -180,3 +180,25 @@ def test_membership_with_duplicate_context_ids(self, monkeypatch):
         monkeypatch.setattr(models.LtiLaunch, "deployment", new_deployment)
         lti_launch = models.LtiLaunch(None)
         assert lti_launch.membership == new_membership
+
+    def test_membership_with_aud_array(self, monkeypatch):
+        membership = factories.LtiMembershipFactory()
+        deployment = membership.context.deployment
+        registration = deployment.registration
+        launch_data = {
+            "iss": registration.issuer,
+            "aud": [registration.client_id],
+            "sub": membership.user.sub,
+            "https://purl.imsglobal.org/spec/lti/claim/deployment_id": (
+                deployment.deployment_id
+            ),
+            "https://purl.imsglobal.org/spec/lti/claim/context": {
+                "id": membership.context.id_on_platform,
+            },
+        }
+        monkeypatch.setattr(
+            models.LtiLaunch, "get_launch_data", lambda self: launch_data
+        )
+        monkeypatch.setattr(models.LtiLaunch, "deployment", deployment)
+        lti_launch = models.LtiLaunch(None)
+        assert lti_launch.membership == membership