module_param(no_msi, bool, 0444);
 MODULE_PARM_DESC(no_msi, "Disable MSI interrupt in pci_endpoint_test");
 
+static int irq_type = IRQ_TYPE_MSI;
+module_param(irq_type, int, 0444);
+MODULE_PARM_DESC(irq_type, "IRQ mode selection in pci_endpoint_test (0 - Legacy, 1 - MSI)");
+
 enum pci_barno {
        BAR_0,
        BAR_1,
 struct pci_endpoint_test_data {
        enum pci_barno test_reg_bar;
        size_t alignment;
-       bool no_msi;
+       int irq_type;
 };
 
 static inline u32 pci_endpoint_test_readl(struct pci_endpoint_test *test,
        pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_SIZE,
                                 size);
 
-       pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_IRQ_TYPE,
-                                no_msi ? IRQ_TYPE_LEGACY : IRQ_TYPE_MSI);
+       pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_IRQ_TYPE, irq_type);
        pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_IRQ_NUMBER, 1);
        pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_COMMAND,
                                 COMMAND_COPY);
 
        pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_SIZE, size);
 
-       pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_IRQ_TYPE,
-                                no_msi ? IRQ_TYPE_LEGACY : IRQ_TYPE_MSI);
+       pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_IRQ_TYPE, irq_type);
        pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_IRQ_NUMBER, 1);
        pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_COMMAND,
                                 COMMAND_READ);
 
        pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_SIZE, size);
 
-       pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_IRQ_TYPE,
-                                no_msi ? IRQ_TYPE_LEGACY : IRQ_TYPE_MSI);
+       pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_IRQ_TYPE, irq_type);
        pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_IRQ_NUMBER, 1);
        pci_endpoint_test_writel(test, PCI_ENDPOINT_TEST_COMMAND,
                                 COMMAND_WRITE);
        test->alignment = 0;
        test->pdev = pdev;
 
+       if (no_msi)
+               irq_type = IRQ_TYPE_LEGACY;
+
        data = (struct pci_endpoint_test_data *)ent->driver_data;
        if (data) {
                test_reg_bar = data->test_reg_bar;
                test->alignment = data->alignment;
-               no_msi = data->no_msi;
+               irq_type = data->irq_type;
        }
 
        init_completion(&test->irq_raised);
 
        pci_set_master(pdev);
 
-       if (!no_msi) {
+       switch (irq_type) {
+       case IRQ_TYPE_LEGACY:
+               break;
+       case IRQ_TYPE_MSI:
                irq = pci_alloc_irq_vectors(pdev, 1, 32, PCI_IRQ_MSI);
                if (irq < 0)
                        dev_err(dev, "Failed to get MSI interrupts\n");
                test->num_irqs = irq;
+               break;
+       default:
+               dev_err(dev, "Invalid IRQ type selected\n");
        }
 
        err = devm_request_irq(dev, pdev->irq, pci_endpoint_test_irqhandler,