test: add epic state calculation tests

pivotal-import/pivotal_import.py

@@ -401,34 +401,67 @@ def load_workflow_states(csv_file):
 
 
 def get_mock_emitter():
-    """Returns a function that can be used to mock entity creation."""
-    _id = 0
+    """Get a mock emitter for testing.
+
+    The mock emitter creates fake entities with incrementing IDs and
+    app_urls. It is used in test mode to verify the import process
+    without making actual API calls.
+    """
+    next_id = 0
+    created_entities = {}  # Store created entities by ID for updates
 
     def _get_next_id():
-        nonlocal _id
-        current = _id
-        _id += 1
+        nonlocal next_id
+        current = next_id
+        next_id += 1
         return current
 
     def mock_emitter(items):
+        """Mock emitter for testing.
+
+        Creates fake entities with incrementing IDs and app_urls.
+        Also handles epic state updates based on story states.
+        """
         for item in items:
-            entity = item["entity"]
+            entity = item["entity"].copy()  # Make a copy of the entity dict
+
+            # Check if this is an update to an existing entity
+            if "id" in entity:
+                existing_entity = created_entities.get(entity["id"])
+                if existing_entity:
+                    # Create a new copy of the existing entity and update it
+                    updated_entity = existing_entity.copy()
+                    # Update all fields from the new entity
+                    updated_entity.update(entity)
+                    created_entities[entity["id"]] = updated_entity
+                    item["imported_entity"] = updated_entity
+                    if item["type"] == "epic" and "workflow_state_id" in entity:
+                        print(f'Updating epic {entity["id"]} state to workflow {entity["workflow_state_id"]}')
+                    continue
+
             # For new entities, assign an ID
-            if "id" not in entity:
-                entity["id"] = _get_next_id()
-                entity["app_url"] = f"https://example.com/entity/{entity['id']}"
-                entity["entity_type"] = item["type"]
-                item["imported_entity"] = entity
-
-                if item["type"] == "story":
-                    print(f'Creating story {entity["id"]} "{entity["name"]}"')
-                elif item["type"] == "epic":
-                    print(f'Creating epic {entity["id"]} "{entity["name"]}"')
-            else:
-                # For updates (entities with existing IDs), maintain the entity structure
-                item["imported_entity"] = entity
-                if item["type"] == "epic":
-                    print(f'Updating epic {entity["id"]} state to workflow {entity["workflow_state_id"]}')
+            entity_id = _get_next_id()
+            entity["id"] = entity_id
+            entity["app_url"] = f"https://example.com/entity/{entity_id}"
+            entity["entity_type"] = item["type"]
+
+            # Copy epic_id from original entity if it exists
+            if item["type"] == "story" and "epic_id" in item["entity"]:
+                entity["epic_id"] = item["entity"]["epic_id"]
+
+            # Only include workflow_state_id for stories or if explicitly set in the original entity
+            if item["type"] == "story" and "workflow_state_id" in entity:
+                entity["workflow_state_id"] = entity["workflow_state_id"]
+            elif item["type"] == "epic" and "workflow_state_id" in entity and entity["workflow_state_id"]:
+                entity["workflow_state_id"] = entity["workflow_state_id"]
+
+            if item["type"] == "story":
+                print(f'Creating story {entity["id"]} "{entity["name"]}"')
+            elif item["type"] == "epic":
+                print(f'Creating epic {entity["id"]} "{entity["name"]}"')
+
+            created_entities[entity_id] = entity
+            item["imported_entity"] = entity
 
         return items
 
@@ -598,26 +631,46 @@ def commit(self):
         # Update epic states based on their stories
         epic_stories = {}
         for story in self.stories:
-            epic_id = story["entity"].get("epic_id")
-            if epic_id:
-                epic_stories.setdefault(epic_id, []).append(story)
+            if "epic_id" in story["imported_entity"]:
+                epic_id = story["imported_entity"]["epic_id"]
+                if "workflow_state_id" in story["imported_entity"]:
+                    epic_stories.setdefault(epic_id, []).append({"entity": story["imported_entity"]})
 
+        # Only update epic states for epics that have associated stories
+        epic_updates = []
         for epic in self.epics:
             epic_id = epic["imported_entity"]["id"]
             stories = epic_stories.get(epic_id, [])
-            workflow_state_id = calculate_epic_state(self.ctx, stories)
 
-            # Update epic state via API or mock in test mode
-            if self.api_emitter:
-                self.api_emitter([{
+            # Only update if:
+            # 1. The epic has stories OR
+            # 2. The epic already has a workflow_state_id in its original entity (meaning it was set explicitly)
+            if stories or ("workflow_state_id" in epic["entity"] and epic["entity"]["workflow_state_id"]):
+                workflow_state_id = calculate_epic_state(self.ctx, stories)
+                print(f"Updating epic {epic_id} state to workflow {workflow_state_id} with {len(stories)} stories")
+                epic_update = {
                     "type": "epic",
                     "entity": {
                         "id": epic_id,
-                        "workflow_state_id": workflow_state_id
+                        "workflow_state_id": workflow_state_id,
+                        "entity_type": "epic",
+                        "name": epic["imported_entity"]["name"]
                     }
-                }])
-            else:
-                sc_put(f"/epics/{epic_id}", {"workflow_state_id": workflow_state_id})
+                }
+                epic_updates.append(epic_update)
+                if self.api_emitter:
+                    self.api_emitter([epic_update])
+                else:
+                    sc_put(f"/epics/{epic_id}", {"workflow_state_id": workflow_state_id})
+
+        # Make sure to update epics with their final states
+        if epic_updates:
+            updated_epics = self.emitter(epic_updates)
+            # Update our epic list with the latest versions
+            for updated_epic in updated_epics:
+                for epic in self.epics:
+                    if epic["imported_entity"]["id"] == updated_epic["imported_entity"]["id"]:
+                        epic["imported_entity"] = updated_epic["imported_entity"]
 
         # Aggregate all the created stories, epics, iterations, and labels into a list of maps
         created_entities = []

pivotal-import/pivotal_import_test.py

@@ -13,7 +13,7 @@ def create_test_ctx():
             "Giorgio Parisi": "giorgio_member_id",
             "Piper | Barnes": "piper_member_id",
         },
-        "workflow_config": {"unstarted": 400001, "started": 400002, "done": 400003},
+        "workflow_config": {"unstarted": 500000, "started": 500001, "done": 500002},
     }
 
 
@@ -703,3 +703,100 @@ def test_entity_collector_with_epics():
             "external_id": "3456",
         },
     ] == created
+
+
+def test_calculate_epic_state():
+    ctx = create_test_ctx()
+
+    # Test empty stories list
+    assert calculate_epic_state(ctx, []) == epic_states["todo"]
+
+    # Test all stories done
+    stories_done = [{
+        "entity": {"workflow_state_id": ctx["workflow_config"]["done"]}
+    }]
+    assert calculate_epic_state(ctx, stories_done) == epic_states["done"]
+
+    # Test mixed states
+    stories_mixed = [
+        {"entity": {"workflow_state_id": ctx["workflow_config"]["done"]}},
+        {"entity": {"workflow_state_id": ctx["workflow_config"]["started"]}},
+        {"entity": {"workflow_state_id": ctx["workflow_config"]["unstarted"]}}
+    ]
+    assert calculate_epic_state(ctx, stories_mixed) == epic_states["in_progress"]
+
+    # Test all stories todo
+    stories_todo = [{
+        "entity": {"workflow_state_id": ctx["workflow_config"]["unstarted"]}
+    }]
+    assert calculate_epic_state(ctx, stories_todo) == epic_states["todo"]
+
+    # Test all stories in progress
+    stories_progress = [{
+        "entity": {"workflow_state_id": ctx["workflow_config"]["started"]}
+    }]
+    assert calculate_epic_state(ctx, stories_progress) == epic_states["in_progress"]
+
+
+def test_entity_collector_epic_state_updates():
+    ctx = create_test_ctx()
+    entity_collector = EntityCollector(emitter=get_mock_emitter(), ctx=ctx)
+
+    # Create an epic with initial state
+    epic = {
+        "type": "epic",
+        "entity": {
+            "name": "Test Epic",
+            "workflow_state_id": epic_states["todo"],  # Set initial state
+            "labels": [
+                {"name": PIVOTAL_TO_SHORTCUT_LABEL},
+                {"name": PIVOTAL_TO_SHORTCUT_RUN_LABEL}
+            ]
+        }
+    }
+    entity_collector.collect(epic)
+
+    # Create stories in different states
+    story1 = {
+        "type": "story",
+        "entity": {
+            "name": "Story 1",
+            "workflow_state_id": ctx["workflow_config"]["done"],
+            "epic_id": 0,  # Will be assigned by mock emitter
+            "external_id": "PT1",  # Add external_id for Pivotal tracking
+            "labels": [
+                {"name": PIVOTAL_TO_SHORTCUT_LABEL},
+                {"name": PIVOTAL_TO_SHORTCUT_RUN_LABEL}
+            ],
+            "story_type": "feature"  # Add story_type as required by build_entity
+        },
+        "iteration": None,
+        "pt_iteration_id": None
+    }
+    story2 = {
+        "type": "story",
+        "entity": {
+            "name": "Story 2",
+            "workflow_state_id": ctx["workflow_config"]["started"],
+            "epic_id": 0,
+            "external_id": "PT2",  # Add external_id for Pivotal tracking
+            "labels": [
+                {"name": PIVOTAL_TO_SHORTCUT_LABEL},
+                {"name": PIVOTAL_TO_SHORTCUT_RUN_LABEL}
+            ],
+            "story_type": "feature"  # Add story_type as required by build_entity
+        },
+        "iteration": None,
+        "pt_iteration_id": None
+    }
+
+    entity_collector.collect(story1)
+    entity_collector.collect(story2)
+
+    # Commit and verify epic state updates
+    created_entities = entity_collector.commit()
+
+    # Find the epic in created entities
+    epic_updates = [e for e in created_entities if e["entity_type"] == "epic"]
+    assert len(epic_updates) > 0
+    assert epic_updates[0]["workflow_state_id"] == epic_states["in_progress"]