diff --git a/.github/workflows/test_upb.yml b/.github/workflows/test_upb.yml
index d302edd0e919..2edc19d22917 100644
--- a/.github/workflows/test_upb.yml
+++ b/.github/workflows/test_upb.yml
@@ -264,13 +264,24 @@ jobs:
       - name: Install Protobuf Test Wheel
         if: ${{ !matrix.continuous-only || inputs.continuous-run }}
         run: pip install -vvv --no-index --find-links wheels protobuftests
-      - name: Run the unit tests
+      - name: Run legacy unit tests
         if: ${{ !matrix.continuous-only || inputs.continuous-run }}
+        # Legacy upb tests use a "test wrapper" and get installed under pb_unit_tests.
         run: |
-          TESTS=$(pip show -f protobuftests | grep pb_unit_tests.*py$ | sed 's,/,.,g' | sed 's,\\,.,g' | sed -E 's,.py$,,g')
-          for test in $TESTS; do
+          PB_UNIT_TESTS=$(pip show -f protobuftests | grep pb_unit_tests.*py$ | sed 's,/,.,g' | sed 's,\\,.,g' | sed -E 's,.py$,,g')
+          for test in $PB_UNIT_TESTS; do
             python -m unittest -v $test
           done
+      - name: Run unit tests
+        if: ${{ !matrix.continuous-only || inputs.continuous-run }}
+        # Newer upb tests are in the standard google.protobuf.internal path.
+        # We will eventually make this into a wildcard rule once all tests
+        # have been migrated to be compatible with upb.
+        run: |
+          TESTS=(message_test message_factory_test descriptor_test proto_builder_test)
+          for test in ${TESTS[@]}; do
+            python -m unittest -v google.protobuf.internal.${test}
+          done
 
   test_pure_python_wheels:
     name: Test Pure Python Wheels
diff --git a/python/google/protobuf/internal/message_test.py b/python/google/protobuf/internal/message_test.py
index ddc903b0386b..9981604230b2 100755
--- a/python/google/protobuf/internal/message_test.py
+++ b/python/google/protobuf/internal/message_test.py
@@ -1456,22 +1456,14 @@ def testAssignInvalidEnum(self):
     """Assigning an invalid enum number is not allowed for closed enums."""
     m = unittest_pb2.TestAllTypes()
 
-    # TODO Enable these once upb's behavior is made conformant.
-    if api_implementation.Type() != 'upb':
-      # Can not assign unknown enum to closed enums.
-      with self.assertRaises(ValueError) as _:
-        m.optional_nested_enum = 1234567
-      self.assertRaises(ValueError, m.repeated_nested_enum.append, 1234567)
-      # Assignment is a different code path than append for the C++ impl.
-      m.repeated_nested_enum.append(2)
-      m.repeated_nested_enum[0] = 2
-      with self.assertRaises(ValueError):
-        m.repeated_nested_enum[0] = 123456
-    else:
+    # Can not assign unknown enum to closed enums.
+    with self.assertRaises(ValueError) as _:
       m.optional_nested_enum = 1234567
-      m.repeated_nested_enum.append(1234567)
-      m.repeated_nested_enum.append(2)
-      m.repeated_nested_enum[0] = 2
+    self.assertRaises(ValueError, m.repeated_nested_enum.append, 1234567)
+    # Assignment is a different code path than append for the C++ impl.
+    m.repeated_nested_enum.append(2)
+    m.repeated_nested_enum[0] = 2
+    with self.assertRaises(ValueError):
       m.repeated_nested_enum[0] = 123456
 
     # Unknown enum value can be parsed but is ignored.