diff --git a/betterproto/tests/inputs/to_dict_with_missing_enum/test_to_dict_with_missing_enum.py b/betterproto/tests/inputs/to_dict_with_missing_enum/test_to_dict_with_missing_enum.py
index 7c56f74..bf5f5fa 100644
--- a/betterproto/tests/inputs/to_dict_with_missing_enum/test_to_dict_with_missing_enum.py
+++ b/betterproto/tests/inputs/to_dict_with_missing_enum/test_to_dict_with_missing_enum.py
@@ -1,7 +1,7 @@
-from betterproto.tests.output_betterproto.to_dict_with_missing_enum.to_dict_with_missing_enum import TestMessage
+from betterproto.tests.output_betterproto.to_dict_with_missing_enum.to_dict_with_missing_enum import TestMessage, TestMessageMyEnum
 
 
 def test_message_attributes():
-    assert TestMessage(x=TestMessage.MyEnum.ONE).to_dict()['x'] == "ONE", "MyEnum.ONE is not serialized to 'ONE'"
-    assert TestMessage(x=TestMessage.MyEnum.THREE).to_dict()['x'] == "THREE", "MyEnum.THREE is not serialized to 'THREE'"
-    assert TestMessage(x=TestMessage.MyEnum.FOUR).to_dict()['x'] == "FOUR", "MyEnum.FOUR is not serialized to 'FOUR'"
+    assert TestMessage(x=TestMessageMyEnum.ONE).to_dict()['x'] == "ONE", "MyEnum.ONE is not serialized to 'ONE'"
+    assert TestMessage(x=TestMessageMyEnum.THREE).to_dict()['x'] == "THREE", "MyEnum.THREE is not serialized to 'THREE'"
+    assert TestMessage(x=TestMessageMyEnum.FOUR).to_dict()['x'] == "FOUR", "MyEnum.FOUR is not serialized to 'FOUR'"