From f764d93a87d31849ac7a5998dfa03f4d65ef5bd2 Mon Sep 17 00:00:00 2001
From: Hariharan Devarajan <mani.hariharan@gmail.com>
Date: Tue, 27 Aug 2024 14:05:59 -0700
Subject: [PATCH] Fixed mocking for DFTracer (#220)

* Fixed mocking for DFTracer

The current mocking was not following the correct signature as used in DFTracer

* Fixed typo

* Update utility.py

Make it more generic

---------

Co-authored-by: Huihuo Zheng <zhenghh04@gmail.com>
---
 dlio_benchmark/utils/utility.py | 41 ++++++++++++++++++++++-----------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/dlio_benchmark/utils/utility.py b/dlio_benchmark/utils/utility.py
index 6bbde3e8..f837ee18 100644
--- a/dlio_benchmark/utils/utility.py
+++ b/dlio_benchmark/utils/utility.py
@@ -37,27 +37,42 @@
     from dftracer.logger import dftracer as PerfTrace, dft_fn as Profile, DFTRACER_ENABLE as DFTRACER_ENABLE
 except:
     class Profile(object):
-        def __init__(self, name=None, cat=None):
-            self.type = type
-        def log(self, func):
-            return func
-        def log_init(self, func):
-            return func
-        def iter(self, a):
-            return a
+        def __init__(self,  **kwargs):
+            return 
+        def log(self,  **kwargs):
+            return 
+        def log_init(self,  **kwargs):
+            return 
+        def iter(self,  **kwargs):
+            return 
         def __enter__(self):
             return
-        def __exit__(self, type, value, traceback):
+        def __exit__(self, **kwargs):
             return
-        def update(self, *, epoch=0, step=0, size=0, default=None):
+        def update(self, **kwargs):
+            return
+        def flush(self):
+            return
+        def reset(self):
+            return
+        def log_static(self, **kwargs):
             return
     class dftracer(object):
         def __init__(self,):
             self.type = None
-        def initialize_log(self, logfile=None, data_dir=None, process_id=-1):
+        def initialize_log(self, **kwargs):
+            return
+        def get_time(self):
+            return
+        def enter_event(self):
+            return
+        def exit_event(self):
+            return
+        def log_event(self,  **kwargs):
+            return
+        def finalize(self):
             return
-        def iter(self, a):
-            return a
+        
     PerfTrace = dftracer()
     DFTRACER_ENABLE = False