gr/cmds/cicd/
mermaid.rs

1use std::{
2    collections::HashMap,
3    fmt::{self, Display, Formatter},
4    ops::Index,
5};
6
7use crate::{error::GRError, Result};
8
9/// A .gitlab-ci.yml is a sequence of stages, where each stage is a collection
10/// of jobs. A stage name is unique, so we can uniquely identify them by name.
11#[derive(Debug)]
12pub struct Stage {
13    pub name: String,
14    pub jobs: Vec<Job>,
15}
16
17impl Stage {
18    pub fn new(name: &str) -> Self {
19        Self {
20            name: name.to_string(),
21            jobs: vec![],
22        }
23    }
24}
25
26type StageName = String;
27
28// Map of stage names to their respective stages.
29#[derive(Debug)]
30pub struct StageMap {
31    stage_names: Vec<StageName>,
32    stages: HashMap<StageName, Stage>,
33}
34
35impl StageMap {
36    fn new() -> Self {
37        Self {
38            stage_names: vec![],
39            stages: HashMap::new(),
40        }
41    }
42
43    fn insert(&mut self, name: StageName, stage: Stage) {
44        self.stage_names.push(name.clone());
45        self.stages.insert(name, stage);
46    }
47
48    fn get_mut(&mut self, name: &str) -> Option<&mut Stage> {
49        self.stages.get_mut(name)
50    }
51
52    fn contains_key(&self, name: &str) -> bool {
53        self.stages.contains_key(name)
54    }
55}
56
57/// A job is a unique unit of work that is executed in a gitlab-ci pipeline. They
58/// belong to a stage. No job can be named the same, so we can uniquely identify
59/// them by name.
60#[derive(Debug)]
61pub struct Job {
62    pub name: String,
63    pub rules: Vec<HashMap<String, CicdEntity>>,
64}
65
66impl Job {
67    pub fn new(name: &str, rules: Vec<HashMap<String, CicdEntity>>) -> Self {
68        Self {
69            name: name.to_string(),
70            rules,
71        }
72    }
73}
74
75/// Defines a CicdEntity entity that can be a sequence, a mapping, a string or null.
76#[derive(Clone, Debug, PartialEq, Eq)]
77pub enum CicdEntity {
78    Vec(Vec<CicdEntity>),
79    Hash(HashMap<String, CicdEntity>),
80    String(String),
81    Integer(i64),
82    Null,
83}
84
85impl CicdEntity {
86    pub fn as_vec(&self) -> Option<&Vec<CicdEntity>> {
87        if let CicdEntity::Vec(ref v) = *self {
88            Some(v)
89        } else {
90            None
91        }
92    }
93
94    pub fn as_hash(&self) -> Option<&HashMap<String, CicdEntity>> {
95        if let CicdEntity::Hash(ref h) = *self {
96            Some(h)
97        } else {
98            None
99        }
100    }
101
102    pub fn as_str(&self) -> Option<&str> {
103        if let CicdEntity::String(ref s) = *self {
104            Some(s)
105        } else {
106            None
107        }
108    }
109}
110
111// Index operations to have similar ergonomics as Yaml.
112impl Index<&str> for CicdEntity {
113    type Output = CicdEntity;
114
115    fn index(&self, index: &str) -> &Self::Output {
116        if let CicdEntity::Hash(ref h) = *self {
117            h.get(index).unwrap_or(&CicdEntity::Null)
118        } else {
119            &CicdEntity::Null
120        }
121    }
122}
123
124impl Index<usize> for CicdEntity {
125    type Output = CicdEntity;
126
127    fn index(&self, index: usize) -> &Self::Output {
128        if let CicdEntity::Vec(ref v) = *self {
129            v.get(index).unwrap_or(&CicdEntity::Null)
130        } else {
131            &CicdEntity::Null
132        }
133    }
134}
135
136pub enum EntityName {
137    Stage,
138    Job,
139}
140
141impl AsRef<str> for EntityName {
142    fn as_ref(&self) -> &str {
143        match self {
144            EntityName::Stage => "stages",
145            EntityName::Job => "jobs",
146        }
147    }
148}
149
150pub trait ToCicdEntity {
151    fn get(&self, entity_name: &Option<EntityName>) -> CicdEntity;
152}
153
154pub trait CicdParser {
155    fn get_stages(&self) -> Result<StageMap>;
156    /// Gathers the jobs and populate the stages with their corresponding jobs
157    fn get_jobs(&self, stages: &mut StageMap);
158}
159
160// Encapsulates the YAML parser library that we use to parse the YAML file.
161pub struct YamlParser<T> {
162    parser: T,
163}
164
165impl<T> YamlParser<T> {
166    pub fn new(parser: T) -> Self {
167        Self { parser }
168    }
169}
170
171impl<T: ToCicdEntity> CicdParser for YamlParser<T> {
172    fn get_stages(&self) -> Result<StageMap> {
173        let entity = self.parser.get(&Some(EntityName::Stage));
174        if let Some(cicd_stage_names) = entity.as_vec() {
175            let mut stages = StageMap::new();
176            for cicd_stage_name in cicd_stage_names {
177                if let Some(stage_name) = cicd_stage_name.as_str() {
178                    let stage = Stage::new(stage_name);
179                    stages.insert(stage_name.to_string(), stage);
180                }
181            }
182            Ok(stages)
183        } else {
184            Err(GRError::MermaidParsingError("No stages found".to_string()).into())
185        }
186    }
187
188    fn get_jobs(&self, stages: &mut StageMap) {
189        let entity = self.parser.get(&Some(EntityName::Job));
190        if let Some(cicd_job_details) = entity.as_hash() {
191            for (job, job_details) in cicd_job_details {
192                let job_name = job.as_str();
193                //verify if the job has a stage
194                let stage = job_details["stage"].as_str();
195                if stage.is_none() {
196                    // could be an anchor `.template` without an associated stage
197                    continue;
198                }
199                let stage = stage.unwrap();
200                // All jobs need a corresponding stage. If the stage does not
201                // exit for this job, then technically is a wrong configuration.
202                // We skip it as it's not a valid job that can be added to a
203                // stage. User will get an error message when pushing project to
204                // GitLab or when linting the file.
205                if !stages.contains_key(stage) {
206                    // could be an anchor `.template` without an associated stage
207                    continue;
208                }
209                let mut rules: Vec<HashMap<String, CicdEntity>> = job_details["rules"]
210                    .as_vec()
211                    .map(|rules| {
212                        rules
213                            .iter()
214                            .map(|rule| {
215                                if let Some(rule) = rule.as_hash() {
216                                    let mut rule_map = HashMap::new();
217                                    for (key, value) in rule.iter() {
218                                        rule_map.insert(key.clone(), value.clone());
219                                    }
220                                    rule_map
221                                } else if let Some(rule) = rule.as_vec() {
222                                    let mut rule_map = HashMap::new();
223                                    for rule in rule {
224                                        if let Some(rule) = rule.as_hash() {
225                                            for (key, value) in rule {
226                                                let value = value.clone();
227                                                rule_map.insert(key.clone(), value);
228                                            }
229                                        }
230                                    }
231                                    rule_map
232                                } else {
233                                    // empty rules
234                                    HashMap::new()
235                                }
236                            })
237                            .collect()
238                    })
239                    .unwrap_or_default();
240                // if job_name has white spaces join them with a hyphen
241                let job_name = job_name.split_whitespace().collect::<Vec<&str>>().join("-");
242                // if rules is empty, check only rules
243                let only = job_details["only"].as_vec();
244                if only.is_some() {
245                    rules = vec![];
246                    for rule in only.unwrap() {
247                        let mut rule_map = HashMap::new();
248                        rule_map.insert("only".to_string(), rule.clone());
249                        rules.push(rule_map);
250                    }
251                } else {
252                    let refs = job_details["only"]["refs"].as_vec();
253                    if refs.is_some() {
254                        rules = vec![];
255                        for rule in refs.unwrap() {
256                            let mut rule_map = HashMap::new();
257                            rule_map.insert("only".to_string(), rule.clone());
258                            rules.push(rule_map);
259                        }
260                    }
261                }
262                // If job begins with dot, then it's a template.
263                if job_name.starts_with('.') {
264                    continue;
265                }
266                let job = Job::new(&job_name, rules.clone());
267                let mut parallel_jobs = vec![];
268                // check if it's a parallel job
269                if let Some(parallel) = job_details["parallel"].as_hash() {
270                    let matrix = parallel.get(&"matrix".to_string()).unwrap();
271                    let matrix = matrix.as_vec().unwrap();
272                    let mut all_values = vec![];
273                    for matrix_item in matrix {
274                        let partial_values = combine_matrix_values(matrix_item);
275                        all_values.push(partial_values);
276                    }
277                    for val_matrix in all_values {
278                        for val in val_matrix {
279                            parallel_jobs
280                                .push(Job::new(&format!("{job_name}-{val}"), rules.clone()))
281                        }
282                    }
283                }
284                if parallel_jobs.is_empty() {
285                    stages.get_mut(stage).unwrap().jobs.push(job);
286                } else {
287                    for parallel_job in parallel_jobs {
288                        stages.get_mut(stage).unwrap().jobs.push(parallel_job);
289                    }
290                }
291            }
292        }
293    }
294}
295
296fn combine_matrix_values(matrix: &CicdEntity) -> Vec<String> {
297    let map = matrix.as_hash().unwrap();
298    let keys = map.keys().collect::<Vec<&String>>();
299    let mut all_values = vec![];
300    let mut previous_values = vec![];
301    let num_matrix_keys = keys.len();
302    for key in keys {
303        let values = if let Some(values) = map.get(key).unwrap().as_vec() {
304            values
305                .iter()
306                .map(|x| x.as_str().unwrap().to_string())
307                .collect::<Vec<String>>()
308        } else {
309            vec![map.get(key).unwrap().as_str().unwrap().to_string()]
310        };
311        let mut new_values = vec![];
312        for value in values.iter() {
313            if previous_values.is_empty() {
314                new_values.push(value.to_string());
315            } else {
316                for previous_value in previous_values.iter() {
317                    new_values.push(format!("{}-{}", previous_value, value.as_str()));
318                    all_values.push(format!("{}-{}", previous_value, value.as_str()));
319                }
320            }
321        }
322        previous_values = new_values;
323    }
324    if all_values.is_empty() && num_matrix_keys == 1 {
325        // This is the case where there is only one key in the matrix. The
326        // values could be an array, so we need one job per value.
327        all_values = previous_values;
328    }
329    all_values
330}
331
332#[derive(Default)]
333pub struct Mermaid {
334    pub buf: Vec<String>,
335}
336
337impl Mermaid {
338    pub fn new() -> Self {
339        Self::default()
340    }
341
342    pub fn push(&mut self, line: String) {
343        self.buf.push(line);
344    }
345}
346
347impl Display for Mermaid {
348    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
349        for line in self.buf.iter() {
350            writeln!(f, "{line}")?;
351        }
352        Ok(())
353    }
354}
355
356#[derive(Eq, PartialEq, Debug)]
357pub enum ChartType {
358    StagesWithJobs,
359    Jobs,
360    Stages,
361}
362
363/// Generate a Mermaid state diagram with each stage encapsulating all its jobs
364/// and the links in between stages.
365pub fn generate_mermaid_stages_diagram(
366    parser: impl CicdParser,
367    chart_type: ChartType,
368) -> Result<Mermaid> {
369    let mut mermaid = Mermaid::new();
370
371    match chart_type {
372        ChartType::StagesWithJobs => {
373            mermaid.push("stateDiagram-v2".to_string());
374            mermaid.push("    direction LR".to_string());
375        }
376        ChartType::Jobs | ChartType::Stages => {
377            mermaid.push("graph LR".to_string());
378        }
379    }
380
381    let mut stages = parser.get_stages()?;
382
383    parser.get_jobs(&mut stages);
384
385    for (i, stage) in stages.stage_names.iter().enumerate() {
386        let stage_obj = stages.stages.get(stage).unwrap();
387        let jobs = &stage_obj.jobs;
388
389        // Replace - for _ in stage name to avoid mermaid errors
390        let stage_name = stage_obj.name.replace('-', "_");
391
392        // Include .pre and .post stages only if they have jobs
393        if (stage_name == ".pre" || stage_name == ".post") && jobs.is_empty() {
394            continue;
395        }
396
397        if chart_type == ChartType::StagesWithJobs {
398            mermaid.push(format!("    state {}{}", stage_name, "{"));
399            let anchor_name = format!("anchorT{i}");
400            mermaid.push("        direction LR".to_string());
401            mermaid.push(format!("        state \"jobs\" as {anchor_name}"));
402            for job in jobs.iter() {
403                mermaid.push(format!("        state \"{}\" as {}", job.name, anchor_name));
404            }
405            mermaid.push(format!("    {}", "}"));
406        }
407
408        // check all next stages for compatibility. If the first stage after
409        // current one is compatible and the second stage after current one is
410        // also compatible, there should not be a link between the first and the
411        // second.
412        'stages: for next_stage_name in stages.stage_names.iter().skip(i + 1) {
413            let next_stage_obj = stages.stages.get(next_stage_name).unwrap();
414            let next_jobs = &next_stage_obj.jobs;
415
416            // Skip .pre and .post stages if they have no jobs
417            if (next_stage_obj.name == ".pre" || next_stage_obj.name == ".post")
418                && next_jobs.is_empty()
419            {
420                continue;
421            }
422
423            // Replace - for _ in stage name to avoid mermaid errors
424            let next_stage_name = next_stage_obj.name.replace('-', "_");
425
426            // if there's compatibility after first stage, there should not be a
427            // link on the second stage. For jobs, we need to continue looping
428            // till we finish all the next stage jobs and link them up. If the
429            // jobs in next stage are compatible we control that with the
430            // `jobs_first_stage_compatible` variable. Once we finish iterating
431            // over the next stage jobs, then we break next stage loop.
432            let mut jobs_first_stage_compatible = false;
433            for job in jobs.iter() {
434                for next_job in next_jobs.iter() {
435                    if rules_compatible(&job.rules, &next_job.rules) {
436                        match chart_type {
437                            ChartType::StagesWithJobs | ChartType::Stages => {
438                                mermaid.push(format!("    {stage_name} --> {next_stage_name}"));
439                                // break as we know this stage is compatible
440                                break 'stages;
441                            }
442                            ChartType::Jobs => {
443                                jobs_first_stage_compatible = true;
444                                mermaid.push(format!("    {} --> {}", job.name, next_job.name));
445                            }
446                        }
447                    }
448                }
449            }
450            if jobs_first_stage_compatible {
451                break 'stages;
452            }
453        }
454    }
455
456    Ok(mermaid)
457}
458
459fn rules_compatible(
460    rules1: &[HashMap<String, CicdEntity>],
461    rules2: &[HashMap<String, CicdEntity>],
462) -> bool {
463    if rules1.is_empty() || rules2.is_empty() {
464        return true;
465    }
466    for rule1 in rules1.iter() {
467        for rule2 in rules2.iter() {
468            if rule1 == rule2 {
469                return true;
470            }
471        }
472    }
473    false
474}
475
476#[cfg(test)]
477mod tests {
478    use super::*;
479
480    #[test]
481    fn test_cicd_entity_variants() {
482        // Test Vec variant
483        let vec_entity = CicdEntity::Vec(vec![CicdEntity::Null]);
484        assert!(matches!(vec_entity, CicdEntity::Vec(_)));
485
486        // Test Hash variant
487        let mut hash_map = HashMap::new();
488        hash_map.insert(String::from("key"), CicdEntity::Null);
489        let hash_entity = CicdEntity::Hash(hash_map);
490        assert!(matches!(hash_entity, CicdEntity::Hash(_)));
491
492        // Test String variant
493        let string_entity = CicdEntity::String(String::from("value"));
494        assert!(matches!(string_entity, CicdEntity::String(_)));
495
496        // Test Integer variant
497        let integer_entity = CicdEntity::Integer(42);
498        assert!(matches!(integer_entity, CicdEntity::Integer(_)));
499
500        // Test Null variant
501        let null_entity = CicdEntity::Null;
502        assert!(matches!(null_entity, CicdEntity::Null));
503    }
504
505    #[test]
506    fn test_as_vec() {
507        let vec_entity = CicdEntity::Vec(vec![CicdEntity::Null]);
508        assert!(vec_entity.as_vec().is_some());
509
510        let string_entity = CicdEntity::String(String::from("value"));
511        assert!(string_entity.as_vec().is_none());
512    }
513
514    #[test]
515    fn test_as_hash() {
516        let mut hash_map = HashMap::new();
517        hash_map.insert(String::from("key"), CicdEntity::Null);
518        let hash_entity = CicdEntity::Hash(hash_map);
519        assert!(hash_entity.as_hash().is_some());
520
521        let string_entity = CicdEntity::String(String::from("value"));
522        assert!(string_entity.as_hash().is_none());
523    }
524
525    #[test]
526    fn test_as_str() {
527        let string_entity = CicdEntity::String(String::from("value"));
528        assert_eq!(string_entity.as_str(), Some("value"));
529
530        let integer_entity = CicdEntity::Integer(42);
531        assert!(integer_entity.as_str().is_none());
532    }
533
534    #[test]
535    fn test_index_str() {
536        let mut hash_map = HashMap::new();
537        hash_map.insert(
538            String::from("key"),
539            CicdEntity::String(String::from("value")),
540        );
541        let hash_entity = CicdEntity::Hash(hash_map);
542
543        assert_eq!(
544            hash_entity["key"],
545            CicdEntity::String(String::from("value"))
546        );
547        assert_eq!(hash_entity["missing"], CicdEntity::Null);
548    }
549
550    #[test]
551    fn test_index_usize() {
552        let vec_entity = CicdEntity::Vec(vec![CicdEntity::String(String::from("value"))]);
553
554        assert_eq!(vec_entity[0], CicdEntity::String(String::from("value")));
555        assert_eq!(vec_entity[1], CicdEntity::Null);
556    }
557
558    use std::collections::HashSet;
559
560    #[derive(Clone)]
561    struct MockCicdEntity {
562        stages: Vec<String>,
563        jobs: HashMap<String, CicdEntity>,
564    }
565
566    impl MockCicdEntity {
567        fn new(stages: Vec<String>, jobs: HashMap<String, CicdEntity>) -> Self {
568            Self { stages, jobs }
569        }
570    }
571
572    impl ToCicdEntity for MockCicdEntity {
573        fn get(&self, entity_name: &Option<EntityName>) -> CicdEntity {
574            match entity_name {
575                Some(EntityName::Stage) => CicdEntity::Vec(
576                    self.stages
577                        .iter()
578                        .map(|s| CicdEntity::String(s.clone()))
579                        .collect(),
580                ),
581                Some(EntityName::Job) => CicdEntity::Hash(self.jobs.clone()),
582                None => CicdEntity::Null,
583            }
584        }
585    }
586
587    // Helper function to create a MockCicdEntity
588    fn create_mock_cicd_entity(stages: Vec<&str>, jobs: Vec<(&str, CicdEntity)>) -> MockCicdEntity {
589        MockCicdEntity::new(
590            stages.into_iter().map(String::from).collect(),
591            jobs.into_iter().map(|(k, v)| (k.to_string(), v)).collect(),
592        )
593    }
594
595    #[test]
596    fn test_parse_simple_job() {
597        let mock = create_mock_cicd_entity(
598            vec!["build"],
599            vec![(
600                "build_job",
601                CicdEntity::Hash(HashMap::from([
602                    ("stage".to_string(), CicdEntity::String("build".to_string())),
603                    (
604                        "script".to_string(),
605                        CicdEntity::Vec(vec![CicdEntity::String("echo \"Building\"".to_string())]),
606                    ),
607                ])),
608            )],
609        );
610
611        let parser = YamlParser::new(mock);
612        let mut stage_map = StageMap::new();
613        stage_map.insert("build".to_string(), Stage::new("build"));
614
615        parser.get_jobs(&mut stage_map);
616
617        assert_eq!(stage_map.stages["build"].jobs.len(), 1);
618        assert_eq!(stage_map.stages["build"].jobs[0].name, "build_job");
619    }
620
621    #[test]
622    fn test_job_has_non_existing_stage_then_do_not_include() {
623        let mock = create_mock_cicd_entity(
624            vec!["build"],
625            vec![(
626                "build_job",
627                CicdEntity::Hash(HashMap::from([
628                    (
629                        "stage".to_string(),
630                        CicdEntity::String("non_existing".to_string()),
631                    ),
632                    (
633                        "script".to_string(),
634                        CicdEntity::Vec(vec![CicdEntity::String("echo \"Building\"".to_string())]),
635                    ),
636                ])),
637            )],
638        );
639
640        let parser = YamlParser::new(mock);
641        let mut stage_map = StageMap::new();
642        stage_map.insert("build".to_string(), Stage::new("build"));
643
644        parser.get_jobs(&mut stage_map);
645
646        assert_eq!(stage_map.stages["build"].jobs.len(), 0);
647    }
648
649    #[test]
650    fn test_is_template_not_job() {
651        let mock = create_mock_cicd_entity(
652            vec!["build"],
653            vec![(
654                ".build_job_template",
655                CicdEntity::Hash(HashMap::from([
656                    ("stage".to_string(), CicdEntity::String("build".to_string())),
657                    (
658                        "script".to_string(),
659                        CicdEntity::Vec(vec![CicdEntity::String("echo \"Building\"".to_string())]),
660                    ),
661                ])),
662            )],
663        );
664
665        let parser = YamlParser::new(mock);
666        let mut stage_map = StageMap::new();
667        stage_map.insert("build".to_string(), Stage::new("build"));
668
669        parser.get_jobs(&mut stage_map);
670
671        assert_eq!(stage_map.stages["build"].jobs.len(), 0);
672    }
673
674    #[test]
675    fn test_parse_job_with_rules() {
676        let mock = create_mock_cicd_entity(
677            vec!["test"],
678            vec![(
679                "test_job",
680                CicdEntity::Hash(HashMap::from([
681                    ("stage".to_string(), CicdEntity::String("test".to_string())),
682                    (
683                        "script".to_string(),
684                        CicdEntity::Vec(vec![CicdEntity::String("echo \"Testing\"".to_string())]),
685                    ),
686                    (
687                        "rules".to_string(),
688                        CicdEntity::Vec(vec![CicdEntity::Hash(HashMap::from([(
689                            "if".to_string(),
690                            CicdEntity::String("$CI_COMMIT_BRANCH == \"main\"".to_string()),
691                        )]))]),
692                    ),
693                ])),
694            )],
695        );
696
697        let parser = YamlParser::new(mock);
698
699        let mut stage_map = StageMap::new();
700        stage_map.insert("test".to_string(), Stage::new("test"));
701
702        parser.get_jobs(&mut stage_map);
703
704        assert_eq!(stage_map.stages["test"].jobs.len(), 1);
705        assert_eq!(stage_map.stages["test"].jobs[0].name, "test_job");
706        assert_eq!(stage_map.stages["test"].jobs[0].rules.len(), 1);
707        assert!(stage_map.stages["test"].jobs[0].rules[0].contains_key("if"));
708    }
709
710    #[test]
711    fn test_parse_job_with_only_no_refs() {
712        let mock = create_mock_cicd_entity(
713            vec!["deploy"],
714            vec![(
715                "deploy_job",
716                CicdEntity::Hash(HashMap::from([
717                    (
718                        "stage".to_string(),
719                        CicdEntity::String("deploy".to_string()),
720                    ),
721                    (
722                        "script".to_string(),
723                        CicdEntity::Vec(vec![CicdEntity::String("echo \"Deploying\"".to_string())]),
724                    ),
725                    (
726                        "only".to_string(),
727                        CicdEntity::Vec(vec![
728                            CicdEntity::String("main".to_string()),
729                            CicdEntity::String("develop".to_string()),
730                        ]),
731                    ),
732                ])),
733            )],
734        );
735
736        let parser = YamlParser::new(mock);
737        let mut stage_map = StageMap::new();
738        stage_map.insert("deploy".to_string(), Stage::new("deploy"));
739
740        parser.get_jobs(&mut stage_map);
741
742        assert_eq!(stage_map.stages["deploy"].jobs.len(), 1);
743        assert_eq!(stage_map.stages["deploy"].jobs[0].name, "deploy_job");
744        assert_eq!(stage_map.stages["deploy"].jobs[0].rules.len(), 2);
745
746        let rules = &stage_map.stages["deploy"].jobs[0].rules;
747        let main_rule = rules.iter().find(|r| r["only"].as_str() == Some("main"));
748        let develop_rule = rules.iter().find(|r| r["only"].as_str() == Some("develop"));
749
750        assert!(main_rule.is_some(), "Rule for 'main' branch not found");
751        assert!(
752            develop_rule.is_some(),
753            "Rule for release branches not found"
754        );
755    }
756
757    #[test]
758    fn test_parse_job_with_only_with_refs() {
759        let mock = create_mock_cicd_entity(
760            vec!["deploy"],
761            vec![(
762                "deploy_job",
763                CicdEntity::Hash(HashMap::from([
764                    (
765                        "stage".to_string(),
766                        CicdEntity::String("deploy".to_string()),
767                    ),
768                    (
769                        "script".to_string(),
770                        CicdEntity::Vec(vec![CicdEntity::String("echo \"Deploying\"".to_string())]),
771                    ),
772                    (
773                        "only".to_string(),
774                        CicdEntity::Hash(HashMap::from([(
775                            "refs".to_string(),
776                            CicdEntity::Vec(vec![
777                                CicdEntity::String("main".to_string()),
778                                CicdEntity::String("/^release-.*$/".to_string()),
779                            ]),
780                        )])),
781                    ),
782                ])),
783            )],
784        );
785
786        let parser = YamlParser::new(mock);
787        let mut stage_map = StageMap::new();
788        stage_map.insert("deploy".to_string(), Stage::new("deploy"));
789
790        parser.get_jobs(&mut stage_map);
791
792        assert_eq!(stage_map.stages["deploy"].jobs.len(), 1);
793        assert_eq!(stage_map.stages["deploy"].jobs[0].name, "deploy_job");
794        assert_eq!(stage_map.stages["deploy"].jobs[0].rules.len(), 2);
795
796        let rules = &stage_map.stages["deploy"].jobs[0].rules;
797        let main_rule = rules.iter().find(|r| r["only"].as_str() == Some("main"));
798        let release_rule = rules
799            .iter()
800            .find(|r| r["only"].as_str() == Some("/^release-.*$/"));
801
802        assert!(main_rule.is_some(), "Rule for 'main' branch not found");
803        assert!(
804            release_rule.is_some(),
805            "Rule for release branches not found"
806        );
807    }
808
809    #[test]
810    fn test_parse_parallel_jobs() {
811        let mock = create_mock_cicd_entity(
812            vec!["test"],
813            vec![(
814                "parallel_job",
815                CicdEntity::Hash(HashMap::from([
816                    ("stage".to_string(), CicdEntity::String("test".to_string())),
817                    (
818                        "script".to_string(),
819                        CicdEntity::Vec(vec![CicdEntity::String(
820                            "echo \"Testing with $PYTHON_VERSION and $DATABASE\"".to_string(),
821                        )]),
822                    ),
823                    (
824                        "parallel".to_string(),
825                        CicdEntity::Hash(HashMap::from([(
826                            "matrix".to_string(),
827                            CicdEntity::Vec(vec![CicdEntity::Hash(HashMap::from([
828                                (
829                                    "PYTHON_VERSION".to_string(),
830                                    CicdEntity::Vec(vec![
831                                        CicdEntity::String("3.7".to_string()),
832                                        CicdEntity::String("3.8".to_string()),
833                                    ]),
834                                ),
835                                (
836                                    "DATABASE".to_string(),
837                                    CicdEntity::Vec(vec![
838                                        CicdEntity::String("mysql".to_string()),
839                                        CicdEntity::String("postgres".to_string()),
840                                    ]),
841                                ),
842                            ]))]),
843                        )])),
844                    ),
845                ])),
846            )],
847        );
848
849        let parser = YamlParser::new(mock);
850        let mut stage_map = StageMap::new();
851        stage_map.insert("test".to_string(), Stage::new("test"));
852
853        parser.get_jobs(&mut stage_map);
854
855        assert_eq!(stage_map.stages["test"].jobs.len(), 4);
856
857        // Create a set of expected job names
858        let expected_job_names: HashSet<String> = [
859            "parallel_job-3.7-mysql",
860            "parallel_job-3.7-postgres",
861            "parallel_job-3.8-mysql",
862            "parallel_job-3.8-postgres",
863            "parallel_job-mysql-3.7",
864            "parallel_job-mysql-3.8",
865            "parallel_job-postgres-3.7",
866            "parallel_job-postgres-3.8",
867        ]
868        .iter()
869        .map(|&s| s.to_string())
870        .collect();
871
872        // Check that each job name in the stage matches one of the expected names
873        for job in &stage_map.stages["test"].jobs {
874            assert!(
875                expected_job_names.contains(&job.name),
876                "Unexpected job name: {}",
877                job.name
878            );
879        }
880
881        // Check that we have all combinations of Python versions and databases
882        let job_names: HashSet<&String> = stage_map.stages["test"]
883            .jobs
884            .iter()
885            .map(|job| &job.name)
886            .collect();
887        assert!(job_names
888            .iter()
889            .any(|name| name.contains("3.7") && name.contains("mysql")));
890        assert!(job_names
891            .iter()
892            .any(|name| name.contains("3.7") && name.contains("postgres")));
893        assert!(job_names
894            .iter()
895            .any(|name| name.contains("3.8") && name.contains("mysql")));
896        assert!(job_names
897            .iter()
898            .any(|name| name.contains("3.8") && name.contains("postgres")));
899    }
900
901    #[test]
902    fn test_parse_parallel_job_one_element_array() {
903        let mock = create_mock_cicd_entity(
904            vec!["test"],
905            vec![(
906                "parallel_job",
907                CicdEntity::Hash(HashMap::from([
908                    ("stage".to_string(), CicdEntity::String("test".to_string())),
909                    (
910                        "script".to_string(),
911                        CicdEntity::Vec(vec![CicdEntity::String(
912                            "echo \"Testing with $RUST_VERSION\"".to_string(),
913                        )]),
914                    ),
915                    (
916                        "parallel".to_string(),
917                        CicdEntity::Hash(HashMap::from([(
918                            "matrix".to_string(),
919                            CicdEntity::Vec(vec![CicdEntity::Hash(HashMap::from([(
920                                "RUST_VERSION".to_string(),
921                                CicdEntity::Vec(vec![
922                                    CicdEntity::String("1.50".to_string()),
923                                    CicdEntity::String("1.60".to_string()),
924                                ]),
925                            )]))]),
926                        )])),
927                    ),
928                ])),
929            )],
930        );
931
932        let parser = YamlParser::new(mock);
933        let mut stage_map = StageMap::new();
934        stage_map.insert("test".to_string(), Stage::new("test"));
935
936        parser.get_jobs(&mut stage_map);
937
938        assert_eq!(stage_map.stages["test"].jobs.len(), 2);
939
940        // Create a set of expected job names
941        let expected_job_names: HashSet<String> = ["parallel_job-1.50", "parallel_job-1.60"]
942            .iter()
943            .map(|&s| s.to_string())
944            .collect();
945
946        // Check that each job name in the stage matches one of the expected names
947        for job in &stage_map.stages["test"].jobs {
948            assert!(
949                expected_job_names.contains(&job.name),
950                "Unexpected job name: {}",
951                job.name
952            );
953        }
954
955        // Check that we have all combinations of Python versions and databases
956        let job_names: HashSet<&String> = stage_map.stages["test"]
957            .jobs
958            .iter()
959            .map(|job| &job.name)
960            .collect();
961        assert!(job_names.iter().any(|name| name.contains("1.50")));
962        assert!(job_names.iter().any(|name| name.contains("1.60")));
963    }
964
965    #[test]
966    fn test_get_stages() {
967        let mock = create_mock_cicd_entity(
968            vec!["build", "test", "deploy"],
969            vec![], // We don't need job definitions for this test
970        );
971
972        let parser = YamlParser::new(mock);
973        let stage_map = parser.get_stages().unwrap();
974
975        assert_eq!(stage_map.stage_names.len(), 3);
976        assert_eq!(stage_map.stage_names[0], "build");
977        assert_eq!(stage_map.stage_names[1], "test");
978        assert_eq!(stage_map.stage_names[2], "deploy");
979    }
980
981    // Mermaid testing
982
983    struct MockParser {
984        stages: Vec<String>,
985        jobs: HashMap<String, Vec<MockJob>>,
986    }
987
988    struct MockJob {
989        name: String,
990        rules: Vec<HashMap<String, CicdEntity>>,
991    }
992
993    impl CicdParser for MockParser {
994        fn get_stages(&self) -> Result<StageMap> {
995            let mut map = StageMap::new();
996            for stage in &self.stages {
997                map.insert(stage.clone(), Stage::new(stage));
998            }
999            Ok(map)
1000        }
1001
1002        fn get_jobs(&self, stages: &mut StageMap) {
1003            for (stage_name, mock_jobs) in &self.jobs {
1004                if let Some(stage) = stages.get_mut(stage_name) {
1005                    stage.jobs = mock_jobs
1006                        .iter()
1007                        .map(|mock_job| Job::new(&mock_job.name, mock_job.rules.clone()))
1008                        .collect();
1009                }
1010            }
1011        }
1012    }
1013
1014    fn create_mock_parser(
1015        stages: Vec<&str>,
1016        jobs: Vec<(&str, Vec<(&str, Vec<HashMap<String, CicdEntity>>)>)>,
1017    ) -> MockParser {
1018        let stages = stages.into_iter().map(String::from).collect();
1019        let jobs = jobs
1020            .into_iter()
1021            .map(|(stage, job_specs)| {
1022                (
1023                    stage.to_string(),
1024                    job_specs
1025                        .into_iter()
1026                        .map(|(name, rules)| MockJob {
1027                            name: name.to_string(),
1028                            rules,
1029                        })
1030                        .collect(),
1031                )
1032            })
1033            .collect();
1034        MockParser { stages, jobs }
1035    }
1036
1037    #[test]
1038    fn test_simple_pipeline() -> Result<()> {
1039        let parser = create_mock_parser(
1040            vec!["build", "test", "deploy"],
1041            vec![
1042                ("build", vec![("compile", vec![])]),
1043                (
1044                    "test",
1045                    vec![("unit-test", vec![]), ("integration-test", vec![])],
1046                ),
1047                ("deploy", vec![("production", vec![])]),
1048            ],
1049        );
1050
1051        let mermaid = generate_mermaid_stages_diagram(parser, ChartType::StagesWithJobs)?;
1052        let diagram = mermaid.to_string();
1053
1054        assert!(diagram.contains("stateDiagram-v2"));
1055        assert!(diagram.contains("direction LR"));
1056        assert!(diagram.contains("state build{"));
1057        assert!(diagram.contains("state test{"));
1058        assert!(diagram.contains("state deploy{"));
1059        assert!(diagram.contains("build --> test"));
1060        assert!(diagram.contains("test --> deploy"));
1061        assert!(diagram.contains("state \"compile\" as anchorT0"));
1062        assert!(diagram.contains("state \"unit-test\" as anchorT1"));
1063        assert!(diagram.contains("state \"integration-test\" as anchorT1"));
1064        assert!(diagram.contains("state \"production\" as anchorT2"));
1065
1066        Ok(())
1067    }
1068
1069    #[test]
1070    fn test_pipeline_with_empty_stage() -> Result<()> {
1071        let parser = create_mock_parser(
1072            vec!["build", "test", "deploy"],
1073            vec![
1074                ("build", vec![("compile", vec![])]),
1075                ("test", vec![]),
1076                ("deploy", vec![("production", vec![])]),
1077            ],
1078        );
1079
1080        let mermaid = generate_mermaid_stages_diagram(parser, ChartType::StagesWithJobs)?;
1081        let diagram = mermaid.to_string();
1082
1083        assert!(diagram.contains("build --> deploy"));
1084        assert!(!diagram.contains("test -->"));
1085
1086        Ok(())
1087    }
1088
1089    #[test]
1090    fn test_pipeline_with_rules() -> Result<()> {
1091        let parser = create_mock_parser(
1092            vec!["build", "test", "deploy"],
1093            vec![
1094                (
1095                    "build",
1096                    vec![(
1097                        "compile",
1098                        vec![HashMap::from([(
1099                            "only".to_string(),
1100                            CicdEntity::String("main".to_string()),
1101                        )])],
1102                    )],
1103                ),
1104                (
1105                    "test",
1106                    vec![(
1107                        "unit-test",
1108                        vec![HashMap::from([(
1109                            "only".to_string(),
1110                            CicdEntity::String("main".to_string()),
1111                        )])],
1112                    )],
1113                ),
1114                (
1115                    "deploy",
1116                    vec![(
1117                        "production",
1118                        vec![HashMap::from([(
1119                            "only".to_string(),
1120                            CicdEntity::String("tags".to_string()),
1121                        )])],
1122                    )],
1123                ),
1124            ],
1125        );
1126
1127        let mermaid = generate_mermaid_stages_diagram(parser, ChartType::StagesWithJobs)?;
1128        let diagram = mermaid.to_string();
1129
1130        assert!(diagram.contains("build --> test"));
1131        assert!(!diagram.contains("test --> deploy"));
1132
1133        Ok(())
1134    }
1135
1136    #[test]
1137    fn test_pipeline_with_pre_and_post_stages() -> Result<()> {
1138        let parser = create_mock_parser(
1139            vec![".pre", "build", "test", ".post"],
1140            vec![
1141                (".pre", vec![("setup", vec![])]),
1142                ("build", vec![("compile", vec![])]),
1143                ("test", vec![("unit-test", vec![])]),
1144                (".post", vec![("cleanup", vec![])]),
1145            ],
1146        );
1147
1148        let mermaid = generate_mermaid_stages_diagram(parser, ChartType::StagesWithJobs)?;
1149        let diagram = mermaid.to_string();
1150
1151        assert!(diagram.contains("state .pre{"));
1152        assert!(diagram.contains("state .post{"));
1153        assert!(diagram.contains(".pre --> build"));
1154        assert!(diagram.contains("test --> .post"));
1155        assert!(diagram.contains("state build{"));
1156        assert!(diagram.contains("state test{"));
1157        assert!(diagram.contains("state \"setup\" as anchorT0"));
1158        assert!(diagram.contains("state \"cleanup\" as anchorT3"));
1159
1160        Ok(())
1161    }
1162
1163    #[test]
1164    fn test_pipeline_with_empty_pre_and_post_stages() -> Result<()> {
1165        let parser = create_mock_parser(
1166            vec![".pre", "build", "test", ".post"],
1167            vec![
1168                (".pre", vec![]),
1169                ("build", vec![("compile", vec![])]),
1170                ("test", vec![("unit-test", vec![])]),
1171                (".post", vec![]),
1172            ],
1173        );
1174
1175        let mermaid = generate_mermaid_stages_diagram(parser, ChartType::StagesWithJobs)?;
1176        let diagram = mermaid.to_string();
1177
1178        assert!(!diagram.contains("state .pre{"));
1179        assert!(!diagram.contains("state .post{"));
1180        assert!(diagram.contains("build --> test"));
1181        assert!(diagram.contains("state build{"));
1182        assert!(diagram.contains("state test{"));
1183
1184        Ok(())
1185    }
1186
1187    #[test]
1188    fn test_pipeline_with_long_names() -> Result<()> {
1189        let parser = create_mock_parser(
1190            vec!["build-and-compile", "run-all-tests", "deploy-to-production"],
1191            vec![
1192                ("build-and-compile", vec![("compile-source-code", vec![])]),
1193                (
1194                    "run-all-tests",
1195                    vec![
1196                        ("run-unit-tests", vec![]),
1197                        ("run-integration-tests", vec![]),
1198                    ],
1199                ),
1200                (
1201                    "deploy-to-production",
1202                    vec![("deploy-to-prod-servers", vec![])],
1203                ),
1204            ],
1205        );
1206
1207        let mermaid = generate_mermaid_stages_diagram(parser, ChartType::StagesWithJobs)?;
1208        let diagram = mermaid.to_string();
1209
1210        assert!(diagram.contains("state build_and_compile{"));
1211        assert!(diagram.contains("state run_all_tests{"));
1212        assert!(diagram.contains("state deploy_to_production{"));
1213        assert!(diagram.contains("build_and_compile --> run_all_tests"));
1214        assert!(diagram.contains("run_all_tests --> deploy_to_production"));
1215
1216        Ok(())
1217    }
1218
1219    #[test]
1220    fn test_simple_pipeline_jobs_only() -> Result<()> {
1221        let parser = create_mock_parser(
1222            vec!["build", "test", "deploy"],
1223            vec![
1224                ("build", vec![("compile", vec![])]),
1225                (
1226                    "test",
1227                    vec![("unit-test", vec![]), ("integration-test", vec![])],
1228                ),
1229                ("deploy", vec![("production", vec![])]),
1230            ],
1231        );
1232
1233        let mermaid = generate_mermaid_stages_diagram(parser, ChartType::Jobs)?;
1234        let diagram = mermaid.to_string();
1235
1236        assert!(diagram.contains("graph LR"));
1237        assert!(diagram.contains("compile --> unit-test"));
1238        assert!(diagram.contains("compile --> integration-test"));
1239        assert!(diagram.contains("unit-test --> production"));
1240        assert!(diagram.contains("integration-test --> production"));
1241        assert!(!diagram.contains("state build{"));
1242        assert!(!diagram.contains("state test{"));
1243        assert!(!diagram.contains("state deploy{"));
1244
1245        Ok(())
1246    }
1247
1248    #[test]
1249    fn test_pipeline_with_rules_jobs_only() -> Result<()> {
1250        let parser = create_mock_parser(
1251            vec!["build", "test", "deploy"],
1252            vec![
1253                (
1254                    "build",
1255                    vec![(
1256                        "compile",
1257                        vec![HashMap::from([(
1258                            "only".to_string(),
1259                            CicdEntity::String("main".to_string()),
1260                        )])],
1261                    )],
1262                ),
1263                (
1264                    "test",
1265                    vec![(
1266                        "unit-test",
1267                        vec![HashMap::from([(
1268                            "only".to_string(),
1269                            CicdEntity::String("main".to_string()),
1270                        )])],
1271                    )],
1272                ),
1273                (
1274                    "deploy",
1275                    vec![(
1276                        "production",
1277                        vec![HashMap::from([(
1278                            "only".to_string(),
1279                            CicdEntity::String("tags".to_string()),
1280                        )])],
1281                    )],
1282                ),
1283            ],
1284        );
1285
1286        let mermaid = generate_mermaid_stages_diagram(parser, ChartType::Jobs)?;
1287        let diagram = mermaid.to_string();
1288
1289        assert!(diagram.contains("graph LR"));
1290        assert!(diagram.contains("compile --> unit-test"));
1291        assert!(!diagram.contains("unit-test --> production"));
1292        assert!(!diagram.contains("compile --> production"));
1293
1294        Ok(())
1295    }
1296
1297    #[test]
1298    fn test_pipeline_with_multiple_jobs_per_stage_jobs_only() -> Result<()> {
1299        let parser = create_mock_parser(
1300            vec!["build", "test", "deploy"],
1301            vec![
1302                ("build", vec![("compile", vec![]), ("lint", vec![])]),
1303                (
1304                    "test",
1305                    vec![("unit-test", vec![]), ("integration-test", vec![])],
1306                ),
1307                ("deploy", vec![("staging", vec![]), ("production", vec![])]),
1308            ],
1309        );
1310
1311        let mermaid = generate_mermaid_stages_diagram(parser, ChartType::Jobs)?;
1312        let diagram = mermaid.to_string();
1313
1314        assert!(diagram.contains("graph LR"));
1315        assert!(diagram.contains("compile --> unit-test"));
1316        assert!(diagram.contains("compile --> integration-test"));
1317        assert!(diagram.contains("lint --> unit-test"));
1318        assert!(diagram.contains("lint --> integration-test"));
1319        assert!(diagram.contains("unit-test --> staging"));
1320        assert!(diagram.contains("unit-test --> production"));
1321        assert!(diagram.contains("integration-test --> staging"));
1322        assert!(diagram.contains("integration-test --> production"));
1323
1324        Ok(())
1325    }
1326
1327    #[test]
1328    fn test_simple_pipeline_stages_only() -> Result<()> {
1329        let parser = create_mock_parser(
1330            vec!["build", "test", "deploy"],
1331            vec![
1332                ("build", vec![("compile", vec![])]),
1333                (
1334                    "test",
1335                    vec![("unit-test", vec![]), ("integration-test", vec![])],
1336                ),
1337                ("deploy", vec![("production", vec![])]),
1338            ],
1339        );
1340
1341        let mermaid = generate_mermaid_stages_diagram(parser, ChartType::Stages)?;
1342        let diagram = mermaid.to_string();
1343
1344        assert!(diagram.contains("graph LR"));
1345        assert!(diagram.contains("build --> test"));
1346        assert!(diagram.contains("test --> deploy"));
1347        assert!(!diagram.contains("state build{"));
1348        assert!(!diagram.contains("state test{"));
1349        assert!(!diagram.contains("state deploy{"));
1350
1351        Ok(())
1352    }
1353}