fix tests

enterprise/integrations/stripe_service.py

@@ -69,7 +69,13 @@ async def find_or_create_customer_by_user_id(user_id: str) -> dict | None:
 
     # Save the stripe customer in the local db
     with session_maker() as session:
-        session.add(StripeCustomer(org_id=org.id, stripe_customer_id=customer.id))
+        session.add(
+            StripeCustomer(
+                keycloak_user_id=user_id,
+                org_id=org.id,
+                stripe_customer_id=customer.id,
+            )
+        )
         session.commit()
 
     logger.info(

enterprise/tests/unit/test_stripe_service_db.py

@@ -3,6 +3,7 @@ This test file verifies that the stripe_service functions properly use the datab
 to store and retrieve customer IDs.
 """
 
+import uuid
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -13,8 +14,10 @@ from integrations.stripe_service import (
 )
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
+from storage.org import Org
 from storage.stripe_customer import Base as StripeCustomerBase
 from storage.stripe_customer import StripeCustomer
+from storage.user import User
 from storage.user_settings import Base as UserBase
 
 
@@ -31,79 +34,113 @@ def session_maker(engine):
     return sessionmaker(bind=engine)
 
 
+@pytest.fixture
+def test_org_and_user(session_maker):
+    """Create a test org and user for use in tests."""
+    test_user_id = uuid.uuid4()
+    test_org_id = uuid.uuid4()
+
+    with session_maker() as session:
+        # Create org
+        org = Org(id=test_org_id, name='test-org', contact_email='testy@tester.com')
+        session.add(org)
+        session.flush()
+
+        # Create user with current_org_id
+        user = User(id=test_user_id, current_org_id=test_org_id)
+        session.add(user)
+        session.commit()
+
+    return test_user_id, test_org_id
+
+
 @pytest.mark.asyncio
-async def test_find_customer_id_by_user_id_checks_db_first(session_maker):
+async def test_find_customer_id_by_user_id_checks_db_first(
+    session_maker, test_org_and_user
+):
     """Test that find_customer_id_by_user_id checks the database first"""
 
+    test_user_id, test_org_id = test_org_and_user
+
     # Set up the mock for the database query result
     with session_maker() as session:
+        # Create stripe customer
         session.add(
             StripeCustomer(
-                keycloak_user_id='test-user-id',
+                keycloak_user_id=str(test_user_id),
+                org_id=test_org_id,
                 stripe_customer_id='cus_test123',
             )
         )
         session.commit()
 
-    with patch('integrations.stripe_service.session_maker', session_maker):
+    with (
+        patch('integrations.stripe_service.session_maker', session_maker),
+        patch('storage.org_store.session_maker', session_maker),
+    ):
         # Call the function
-        result = await find_customer_id_by_user_id('test-user-id')
+        result = await find_customer_id_by_user_id(str(test_user_id))
 
         # Verify the result
         assert result == 'cus_test123'
 
 
 @pytest.mark.asyncio
-async def test_find_customer_id_by_user_id_falls_back_to_stripe(session_maker):
+async def test_find_customer_id_by_user_id_falls_back_to_stripe(
+    session_maker, test_org_and_user
+):
     """Test that find_customer_id_by_user_id falls back to Stripe if not found in the database"""
 
+    test_user_id, test_org_id = test_org_and_user
+
     # Set up the mock for stripe.Customer.search_async
     mock_customer = stripe.Customer(id='cus_test123')
     mock_search = AsyncMock(return_value=MagicMock(data=[mock_customer]))
 
     with (
         patch('integrations.stripe_service.session_maker', session_maker),
+        patch('storage.org_store.session_maker', session_maker),
         patch('stripe.Customer.search_async', mock_search),
     ):
         # Call the function
-        result = await find_customer_id_by_user_id('test-user-id')
+        result = await find_customer_id_by_user_id(str(test_user_id))
 
         # Verify the result
         assert result == 'cus_test123'
 
-    # Verify that Stripe was searched
+    # Verify that Stripe was searched with the org_id
     mock_search.assert_called_once()
-    assert "metadata['user_id']:'test-user-id'" in mock_search.call_args[1]['query']
+    assert f"metadata['org_id']:'{str(test_org_id)}'" in mock_search.call_args[1]['query']
 
 
 @pytest.mark.asyncio
-async def test_create_customer_stores_id_in_db(session_maker):
+async def test_create_customer_stores_id_in_db(session_maker, test_org_and_user):
     """Test that create_customer stores the customer ID in the database"""
 
-    # Set up the mock for stripe.Customer.search_async
+    test_user_id, test_org_id = test_org_and_user
+
+    # Set up the mock for stripe.Customer.search_async and create_async
     mock_search = AsyncMock(return_value=MagicMock(data=[]))
     mock_create_async = AsyncMock(return_value=stripe.Customer(id='cus_test123'))
 
     with (
         patch('integrations.stripe_service.session_maker', session_maker),
+        patch('storage.org_store.session_maker', session_maker),
         patch('stripe.Customer.search_async', mock_search),
         patch('stripe.Customer.create_async', mock_create_async),
-        patch(
-            'server.auth.token_manager.TokenManager.get_user_info_from_user_id',
-            AsyncMock(return_value={'email': 'testy@tester.com'}),
-        ),
     ):
         # Call the function
-        result = await find_or_create_customer_by_user_id('5594c7b6-f959-4b81-92e9-b09c206f5081')
+        result = await find_or_create_customer_by_user_id(str(test_user_id))
 
     # Verify the result
-    assert result == 'cus_test123'
+    assert result == {'customer_id': 'cus_test123', 'org_id': str(test_org_id)}
 
     # Verify that the stripe customer was stored in the db
     with session_maker() as session:
         customer = session.query(StripeCustomer).first()
         assert customer.id > 0
-        assert customer.keycloak_user_id == 'test-user-id'
+        assert customer.keycloak_user_id == str(test_user_id)
+        assert customer.org_id == test_org_id
         assert customer.stripe_customer_id == 'cus_test123'
         assert customer.created_at is not None
         assert customer.updated_at is not None